from inspect import (
    _empty,
    getargvalues,
    getouterframes,
    signature,
)

from os import getcwd
from re import compile
from sys import settrace, path
from time import perf_counter_ns


lib_paths = (
    '_bootstr',
    '<string>',
)

class Tracer:
    def __init__(self,
        label       = 'tracing',
        full_path   = False,
        show_types  = False,
        check_types = False,
        show_times  = False,
        show_stamp  = False,
        stdlib      = False,
        args_format = 'str',
        dest        = None,
        buffered    = True,
        echo        = True,
        quiet       = True,
    ):
        self.label = label
        self.level = 0
        self.full_path = full_path
        self.traces = list()
        self.show_types = show_types
        self.check_types = check_types
        self.ignored_paths = set()
        self.ignored_funcs = set()
        self.truncated = list()
        self.show_times = show_times
        self.show_stamp = show_stamp
        self.args_format = args_format
        self.dest = dest or getcwd()
        self.buffered = buffered
        self.echo = echo
        self.quiet = quiet
        self.stdlib = stdlib
        self.setup()

    @property
    def output(self):
        return f'{self.dest}/{self.label}.trace'

    def create(self):
        try:
            with open(self.output, 'x'):
                self.log(f'>>> {self.label}', ts=False)
        except FileExistsError: pass

    def write(self, line):
        with open(self.output, 'a') as out:
            out.write(f'{line}\n')

    def log(self, msg, ts=True):
        ts = f'{self.timestamp()}@{self.label[:12]:<12}' if ts and self.show_stamp else ''
        tab = '    ' * self.level
        msg = f'{ts}{tab}{msg}'
        if self.buffered:
            self.traces.append(msg)
        else:
            self.write(msg)
        if self.echo:
            print(msg)

    def log_exc(self, exc):
        if not self.quiet:
            msg = f'! {exc} !'
            self.log(msg)

    def dump(self):
        self.create()
        with open(self.output, 'a') as out:
            out.write('\n'.join(self.traces))

    def truncate(self, path):
        if not self.full_path:
            for pref in self.truncated:
                if path.startswith(pref):
                    path = path[len(pref)+1:]
        return path.removesuffix('.py')

    def ignore_path(self, path):
        return self.ignored_paths.add(compile(path).search)

    def ignored_path(self, path):
        return any(ignore(path) for ignore in self.ignored_paths)

    def ignore_func(self, func):
        return self.ignored_funcs.add(compile(func).search)

    def ignored_func(self, func):
        return any(ignore(func) for ignore in self.ignored_funcs)

    def truncate_path(self, prefix):
        self.truncated.append(prefix)

    def timestamp(self):
        return perf_counter_ns()

    def start_timer(self):
        if self.show_times:
            self.timers[self.level] = self.timestamp()

    def stop_timer(self):
        if self.show_times:
            self.timers[self.level] = self.timestamp() - self.timers[self.level]

    @property
    def elapsed(self):
        if self.show_times:
            elapsed = self.timers[self.level] // 1_000_000
            if elapsed < 10000:
                return f' [{elapsed} ms]' if elapsed > 0 else ''
            else:
                return f' [{elapsed // 1000} s]'
        else:
            return ''

    def setup(self):
        std_paths = [p for p in path if 'python' in p]
        std_paths.extend(lib_paths)
        if not self.stdlib:
            for p in std_paths:
                self.ignore_path(p)
        if not self.full_path:
            self.truncated.extend(std_paths)
        self.ignore_func(r'Tracer\..*')
        self.ignore_path('tracer')
        self.truncate_path(getcwd())
        self._exc_ = False
        self.timers = dict()

    def format_args(self, args):
        if self.args_format == 'repr':
            formatter = repr
        elif self.args_format == 'str':
            formatter = str
        else:
            raise ValueError(f'Invalid args_formatter : {self.args_format}')

        items = list()
        for k in args:
            t = type(args[k]).__name__
            try:
                v = formatter(args[k])
            except Exception as x:
                self.log_exc(x)
                v = ' ! '
            items.append((k, t, v))
        return items

    def trace(self, file, line, func, args):
        if isinstance(args, tuple) and args[0] in ('>>>', '!!!'):
            if args[0] == '>>>':
                val = args[1]
                if self.show_types:
                    args = f' -> {type(val).__name__}: {val}{self.elapsed}'
                else:
                    args = f' -> {val or ""}{self.elapsed}'
            elif args[0] == '!!!':
                args = f' ! {args[1].__name__} ! {args[2]}'
        else:
            if self.show_types:
                args = f"({', '.join(f'{k}:{t} = {v}' for k, t, v in self.format_args(args))})"
            else:
                args = f"({', '.join(f'{k}={v}' for k, t, v in self.format_args(args))})"
        self.log(f'{file}:{line} {func}{args}')

    def ann_err(self, name, expected, actual):
        if name == '_return_':
            msg = f'! "should return "{expected.__name__}" not "{actual.__name__}" !'
        else:
            msg = f'! arg "{name}" should be "{expected.__name__}" not "{actual.__name__}" !'
        self.log(msg)

    def get_func_args(self, frame):
        arginfo = getargvalues(frame)
        return {
            k:arginfo.locals[k]
            for k in arginfo.args
            if k != 'self'
        }

    def get_sig(self, frame):
        fct = frame.f_code.co_name
        if fct.endswith('comp>'): return
        slf = frame.f_locals.get('self')
        if slf:
            # method
            func = getattr(slf, fct)
        else:
            try:
                # global function
                func = frame.f_globals[fct]
            except KeyError:
                # inner function
                outer = getouterframes(frame)[1]
                func = outer.frame.f_locals[fct]
        return signature(func)
    def chk_arg_ann(self, name, arg, ann):
        if ann is _empty: return
        if ann is None and arg is None: return
        try:
            if not isinstance(arg, ann):
                self.ann_err(name, ann, type(arg))
        except TypeError:
            self.log(f'! unable to  check arg ! {arg}:{ann}')

    def check_sig(self, frame, args):
        if self.check_types:
            sig = self.get_sig(frame)
            if not sig or not sig.parameters: return
            if isinstance(args, tuple) and args[0] == '>>>':
                arg = args[1]
                ann = sig.return_annotation
                self.chk_arg_ann('_return_', arg, ann)
            else:
                if isinstance(args, dict):
                    for name, p in sig.parameters.items():
                        arg = args[name]
                        ann = p.annotation
                        if ann is _empty: continue
                        self.chk_arg_ann(name, arg, ann)

    def get_func_name(self, frame):
        fct = frame.f_code.co_name
        slf = frame.f_locals.get('self')
        try:
            repr(slf) # check handlability
        except Exception:
            slf = None
        if slf:
            try:
                func = getattr(slf, fct)
            except Exception:
                func = None
        else:
            func = frame.f_globals.get(fct)
        try:
            return func.__qualname__
        except Exception:
            return fct

    def get_trace(self, frame, arg, lvl=0):
        path = frame.f_code.co_filename
        if self.ignored_path(path): return False
        func = self.get_func_name(frame)
        if self.ignored_func(func): return False
        self.level += lvl
        self.check_sig(frame, arg)
        path = self.truncate(path)
        line = frame.f_lineno
        self.trace(path, line, func, arg)
        return path, line, func, arg

    def trace_exception(self, frame, typ, val, tbk):
        ret = self.get_trace(frame, ('!!!', typ, val))
        if ret:
            self._exc_ = True

    def trace_return(self, frame, val):
        self.stop_timer()
        if self._exc_:
            self.level-=1
            self._exc_ = False
            return
        if self.get_trace(frame, ('>>>', val)):
            self.level-=1

    def trace_call(self, frame, event, arg):
        args = self.get_func_args(frame)
        self.get_trace(frame, args, 1)
        self.start_timer()
        return self.trace_event

    def trace_event(self, frame, event, arg):
        if event == 'call':
            return self.trace_call

        if event == 'return':
            self.trace_return(frame, arg)

        if event == 'exception':
            self.trace_exception(frame, *arg)

    def __enter__(self):
        if not self.buffered:
            self.create()
        self.level = 0
        self.log(f'{self.label} >>>', ts=False)
        settrace(self.trace_event)

    def __exit__(self, typ, val, tbk):
        settrace(None)
        if self.level == 0:
            self.log(f'{self.label} <<<', ts=False)
        else:
            self.log(f'{self.label} !0!', ts=False)
        if self.buffered:
            self.dump()
        return self.quiet
