#!/usr/bin/env python3
# -*- coding: utf-8 -*-

from __future__ import annotations
from intercept import *


FUNCTION_ERRORS: dict[str, list[str]] = {
    'malloc': ['ENOMEM'],
    'calloc': ['ENOMEM'],
    'realloc': ['ENOMEM'],
    'reallocarray': ['ENOMEM'],
    'read': ['EINTR', 'EIO', 'ECONNRESET', 'ENOTCONN', 'ETIMEDOUT'],
    'pread': ['EINTR', 'EIO'],
    'write': ['EINTR', 'EIO', 'EDQUOT', 'EFBIG', 'ENOSPC', 'EPIPE', 'ECONNRESET'],
    'pwrite': ['EINTR', 'EIO', 'EDQUOT', 'EFBIG', 'ENOSPC', 'EPIPE'],
    'close': [],
    'sigaction': [],
    'sem_init': ['ENOSYS'],
    'sem_open': ['EACCES', 'EEXIST', 'EMFILE', 'ENFILE', 'ENOENT', 'ENOMEM'],
    'sem_post': [],
    'sem_wait': ['EINTR'],
    'sem_trywait': ['EINTR'],
    'sem_timedwait': ['EINTR', 'ETIMEDOUT'],
    'sem_getvalule': [],
    'sem_close': [],
    'sem_unlink': ['EACCES', 'ENOENT'],
    'sem_destroy': [],
    'shm_open': ['EACCES', 'EEXIST', 'EMFILE', 'ENFILE', 'ENOENT'],
    'shm_unlink': ['EACCES', 'ENOENT'],
    'mmap': ['EACCES', 'EBADF', 'EMFILE', 'ENODEV', 'ENOMEM', 'ENOTSUP', 'ENXIO'],
    'munmap': [],
    'ftruncate': ['EINTR', 'EFBIG', 'EIO', 'EBADF'],
    'fork': ['ENOMEM', 'ENOSYS'],
    'wait': ['ECHILD', 'EINTR'],
    'waitpid': ['ECHILD', 'EINTR'],
    'execl': ['ECHILD', 'EINTR'],
    'execlp': ['ECHILD', 'EINTR'],
    'execle': ['ECHILD', 'EINTR'],
    'execv': ['ECHILD', 'EINTR'],
    'execvp': ['ECHILD', 'EINTR'],
    'execvpe': ['ECHILD', 'EINTR'],
    'execve': ['ECHILD', 'EINTR'],
    'fexecve': ['ECHILD', 'EINTR'],
    'pipe': ['EMFILE', 'ENFILE'],
    'dup': ['EBADF', 'EMFILE'],
    'dup2': ['EBADF', 'EBUSY', 'EINTR', 'EMFILE'],
    'dup3': ['EBADF', 'EBUSY', 'EINTR', 'EMFILE'],
    'socket': ['EAFNOSUPPORT', 'EMFILE', 'ENFILE', 'ENOMEM', 'ENOBUFS', 'EACCES'],
    'bind': ['EADDRINUSE', 'EADDRNOTAVAIL', 'EAFNOSUPPORT', 'EACCES'],
    'listen': ['EACCES', 'ENOBUFS', 'EADDRINUSE'],
    'accept': ['ECONNABORTED', 'EINTR', 'EPERM', 'EPROTO', 'ENOMEM'],
    'connect': ['EACCES', 'EPERM', 'ECONNREFUSED', 'ENETUNREACH'],
    'getaddrinfo': [],
    'freeaddrinfo': [],
    'send': ['EINTR'],
    'recv': ['EINTR'],
}

SKIP_ERRORS: list[str] = ['EINTR']


class MemoryAllocationParser(Parser):
    allocated: dict[int, tuple[FunctionCallId, int]]
    invalid_frees: list[tuple[FunctionCallId, int]]
    max_allocated: int
    num_alloc: int
    num_realloc: int
    num_free: int

    def before(self):
        self.allocated = {}
        self.invalid_frees = []
        self.max_allocated = 0
        self.num_alloc = 0
        self.num_realloc = 0
        self.num_free = 0

    def update_max_allocated(self):
        total = sum(a[1] for a in self.allocated.values())
        if total > self.max_allocated:
            self.max_allocated = total

    def after_malloc(self, size, ret_value, errno=None) -> None:
        if ret_value != 0:
            self.num_alloc += 1
            self.allocated[ret_value] = (self.get_call_id('malloc'), size)
        self.update_max_allocated()

    def after_calloc(self, nmemb, size, ret_value, errno=None) -> None:
        if ret_value != 0:
            self.num_alloc += 1
            self.allocated[ret_value] = (self.get_call_id('calloc'), nmemb * size)
        self.update_max_allocated()

    def after_realloc(self, ptr, size, ret_value, errno=None) -> None:
        if ptr == 0:
            if ret_value != 0:
                self.num_alloc += 1
                self.allocated[ret_value] = (self.get_call_id('realloc'), size)
        else:
            if ret_value != 0 and ptr in self.allocated:
                self.num_realloc += 1
                v = self.allocated[ptr]
                del self.allocated[ptr]
                self.allocated[ret_value] = (v[0], size)
        self.update_max_allocated()

    def after_reallocarray(self, ptr, nmemb, size, ret_value, errno=None) -> None:
        if ptr == 0:
            if ret_value != 0:
                self.num_alloc += 1
                self.allocated[ret_value] = (self.get_call_id('reallocarray'), nmemb * size)
        else:
            if ret_value != 0:
                self.num_realloc += 1
                v = self.allocated[ptr]
                del self.allocated[ptr]
                self.allocated[ret_value] = (v[0], nmemb * size)
        self.update_max_allocated()

    def after_getaddrinfo(self, node, service, hints, res_ptr, ret_value, errno=None, res=None) -> None:
        if ret_value.raw == 0 and res is not None:
            self.num_alloc += 1
            size = sum(48 + r['ai_addrlen'] for r in res.target)
            self.allocated[res.ptr] = (self.get_call_id('getaddrinfo'), size)
        self.update_max_allocated()

    def after_getline(self, line_ptr, n_ptr, stream, ret_value, errno=None, n=None, line=None) -> None:
        if ret_value >= 0 and n is not None and line is not None:
            if n_ptr.target == 0 or line_ptr.target == 0:
                self.num_alloc += 1
                self.allocated[line.ptr] = (self.get_call_id('getline'), n)
            elif line_ptr.target != line.ptr:
                self.num_realloc += 1
                v = self.allocated[line_ptr.target]
                del self.allocated[line_ptr.target]
                self.allocated[line.ptr] = (v[0], n)
        self.update_max_allocated()

    def after_getdelim(self, line_ptr, n_ptr, delim, stream, ret_value, errno=None, n=None, line=None) -> None:
        if ret_value >= 0 and n is not None and line is not None:
            if line_ptr.target == 0:
                self.num_alloc += 1
                self.allocated[line.ptr] = (self.get_call_id('getdelim'), n)
            elif line_ptr.target != line.ptr:
                self.num_realloc += 1
                v = self.allocated[line.ptr]
                del self.allocated[line.ptr]
                self.allocated[ret_value] = (v[0], n)
        self.update_max_allocated()

    def after_free(self, ptr) -> None:
        self.num_free += 1
        if ptr != 0:
            if ptr in self.allocated:
                del self.allocated[ptr]
            else:
                self.num_free -= 1
                self.invalid_frees.append((self.get_call_id('free'), ptr))

    def after_freeaddrinfo(self, res: Pointer) -> None:
        self.num_free += 1
        if res != 0:
            if res in self.allocated:
                del self.allocated[res]
            else:
                self.num_free -= 1
                self.invalid_frees.append((self.get_call_id('freeaddrinfo'), res))


class MemoryAllocationTester(MemoryAllocationParser, Handler):
    pass


class InterruptedCheckParser(Parser):
    cycles: int = 50
    functions: dict[str, tuple[str or None, str]] = {
        fn: ('fail EINTR' if fn not in ('sem_post',) else None,
             'return 0' if fn.startswith('sem_') else 'ok')
        for fn, errors in FUNCTION_ERRORS.items()
        if 'EINTR' in errors or fn in ('sem_post',)
    }
    context: dict

    counter: int = 0
    last_call: Optional[FunctionCallId] = None
    tested_function_calls: dict[FunctionCallId, str]

    @property
    def while_testing(self) -> bool:
        return self.counter % self.cycles != 0

    def before(self) -> None:
        self.tested_function_calls = {}

    def after(self) -> None:
        if self.while_testing:
            self.error()
        self.context['results'] = self.tested_function_calls

    def error(self):
        self.tested_function_calls[self.last_call] = 'failed'
        self.counter = 0
        self.last_call = None

    def before_fallback(self, func_name: str, *args) -> str:
        call = self.get_call_id(func_name)
        if self.while_testing and self.last_call != call:
            self.error()
            return 'ok'
        elif func_name not in self.functions:
            return 'ok'
        elif self.functions[func_name][0] is None:
            return self.functions[func_name][1]
        self.counter += 1
        if self.while_testing:
            self.last_call = call
            self.tested_function_calls[self.last_call] = 'running'
            return self.functions[func_name][0]
        else:
            self.tested_function_calls[self.last_call] = 'passed'
            self.last_call = None
            return self.functions[func_name][1]


class InterruptedCheckTester(InterruptedCheckParser, Handler):
    pass


class CallTreeNode(NamedTuple):
    call: FunctionCallId
    depth: int
    data: dict
    children: dict[str, CallTreeNode]


class ReturnValueCheckParser(Parser):
    functions: dict[str, list[str]] = {
        fn: [e for e in errors if e not in SKIP_ERRORS]
        for fn, errors in FUNCTION_ERRORS.items()
        if len(set(errors) - set(SKIP_ERRORS)) > 0
    }
    context: dict
    call_depth: int = 0

    @property
    def call_tree(self) -> CallTreeNode:
        return self.context['call_tree']

    @property
    def last_call(self) -> CallTreeNode:
        return self.context['last_call']

    @property
    def last_return(self) -> str:
        return self.context['last_return']

    @property
    def next(self) -> tuple[Optional[CallTreeNode], Optional[str]]:
        return self.context['next']

    @property
    def is_init(self) -> bool:
        return self.context['state'] == 'init'

    @property
    def is_testing(self) -> bool:
        return self.context['state'] == 'testing'

    @property
    def is_waiting(self) -> bool:
        return self.context['state'] == 'waiting'

    @property
    def is_failed(self) -> bool:
        return self.context['state'] == 'failed'

    def get_next(self, last: Optional[tuple[CallTreeNode, Optional[str]]]) -> tuple[Optional[CallTreeNode], Optional[str]]:
        if last is None:
            if self.call_tree is None:
                return None, None
            last_call, last_err = self.call_tree, None
        else:
            last_call, last_err = last
        if last_err is None:
            if len(last_call.children) == 0:
                return None, None
            err = self.functions.get(last_call.children['ok'].call.func_name, [None])[0]
            return (last_call.children['ok'], err) if err else self.get_next((last_call.children['ok'], err))
        errors = self.functions.get(last_call.call.func_name, [])
        i = errors.index(last_err)
        return (last_call, errors[i + 1]) if i + 1 < len(errors) else self.get_next((last_call, None))

    def before(self) -> None:
        self.context['pid'] = self.pid
        if 'state' not in self.context:
            self.context['state'] = 'init'
            self.context['call_tree'] = None
            self.context['next'] = None
        if self.is_waiting or self.is_failed:
            self.context['state'] = 'testing'
        self.context['last_call'] = None
        self.context['last_return'] = None

    def after(self) -> None:
        if self.context['pid'] != self.pid:
            return
        if self.is_init:
            self.context['state'] = 'testing'
            self.context['next'] = self.get_next(None)
        elif self.is_testing and self.next[0] is None:
            self.context['state'] = 'finished'
        if self.last_call and self.last_call.call.func_name != 'exit':
            self.last_call.children[self.last_return] = (CallTreeNode(FunctionCallId.for_exit(), self.last_call.depth + 1, {}, {}))

    def before_fallback(self, func_name: str, *args) -> str:
        if self.context['pid'] != self.pid:
            return 'ok'

        self.call_depth += 1
        call = self.get_call_id(func_name)
        if self.last_call is None:
            if self.call_tree is None:
                self.context['call_tree'] = CallTreeNode(call, 0, {}, {})
            self.context['last_call'] = self.call_tree
        else:
            if self.last_return not in self.last_call.children:
                self.last_call.children[self.last_return] = CallTreeNode(call, self.last_call.depth + 1, {}, {})
            self.context['last_call'] = self.last_call.children[self.last_return]

        if self.is_init:
            self.context['last_return'] = 'ok'
            return 'ok'

        nxt = self.next
        if nxt[0] is None or call != nxt[0].call or self.call_depth - 1 != nxt[0].depth or nxt[1] is None:
            self.context['last_return'] = 'ok'
            return 'ok'

        self.context['next'] = self.get_next(nxt)
        self.context['state'] = 'waiting' if self.context['next'][1] else 'finished'
        self.context['last_return'] = 'fail ' + nxt[1]
        return self.context['last_return']


class ReturnValueCheckTester(ReturnValueCheckParser, Handler):
    pass


def init_with_ctx(parser: type[Parser]) -> tuple[dict, type[Parser]]:
    ctx = {}
    class ContextParser(parser):
        context: dict = ctx
    return ctx, ContextParser
