import contextlib
import linecache
import sys
import inspect
import functools
import threading
import types
from datetime import datetime
from typing import Optional, Union, TextIO, Callable, Set, Any
from pathlib import Path

from .exceptions import NotAsyncFunctionError
from .utils import OutputWriter, open_output_stream, close_output_stream
from .output import OutputHandler, OutputFormatter
from .frame_watcher import AsyncFrameWatcher


_TLocal = threading.local()


class AsyncSnooper:
    """
    Snoop on the async function, writing everything it's doing to stderr.
    """

    def __init__(
        self,
        output: Optional[Union[str, Path, TextIO, Callable[[str], None]]] = None,
        depth: int = 1,
        watch: tuple = (),
        watch_explode: tuple = (),
        prefix: str = "",
        overwrite: bool = False,
        thread_info: bool = False,
        color: bool = True,
    ):
        self.output_spec = output
        self.depth = depth
        self.watch = watch
        self.watch_explode = watch_explode
        self.prefix = prefix
        self.overwrite = overwrite
        self.thread_info = thread_info
        self.color = color
        self.target_code_objects: Set = set()
        self.stream: Optional[Union[TextIO, Callable[[str], None]]] = None
        self.opened_by_us: bool = False

    def __call__(self, func: inspect.CO_COROUTINE):
        if not inspect.iscoroutinefunction(func):
            raise NotAsyncFunctionError(
                f"@{self.__class__.__name__} can only be applied to async functions, "
                f"but {func.__name__} is {type(func).__name__}"
            )

        self.target_code_objects.add(func.__code__)

        @functools.wraps(func)
        async def wrapper(*args, **kwargs):
            return await self._trace_call(func, args, kwargs)

        return wrapper

    async def _trace_call(self, func: inspect.CO_COROUTINE, args, kwargs):
        self.stream = open_output_stream(self.output_spec, self.overwrite)
        self.opened_by_us = self.output_spec is not None and isinstance(self.output_spec, (str, Path))

        try:
            formatter = OutputFormatter(
                prefix=self.prefix,
                include_thread_info=self.thread_info,
                use_colors=self.color,
            )
            writer = OutputWriter(self.stream)
            handler = OutputHandler(writer, formatter)

            watcher = AsyncFrameWatcher(
                handler=handler,
                snooper=self,
                target_code_objects=self.target_code_objects,
                depth=self.depth,
                watch_expressions=self.watch,
                watch_explode_expressions=self.watch_explode,
            )

            old_trace = sys.gettrace()
            sys.settrace(watcher.trace)

            try:
                result = await func(*args, **kwargs)
                return result
            finally:
                sys.settrace(old_trace)

        finally:
            close_output_stream(self.stream, self.opened_by_us)
            self.stream = None


class AsyncSnoopContext:
    """
    Snoop on the async function context if you want not use wrapper.
    """

    def __init__(
        self,
        output: Optional[Union[str, Path, TextIO, Callable[[str], None]]] = None,
        depth: int = 1,
        watch: tuple = (),
        watch_explode: tuple = (),
        prefix: str = "",
        overwrite: bool = False,
        thread_info: bool = False,
        color: bool = True,
    ):
        self.output_spec = output
        self.depth = depth
        self.watch = watch
        self.watch_explode = watch_explode
        self.overwrite = overwrite
        self.thread_info = thread_info
        self.color = color

        self.target_frames: Set = set()
        self.frame_to_local_reprs = {}
        self.start_times = {}
        self.thread_local = threading.local()
        self._handler = None
        self._formatter = None
        if prefix and not prefix.endswith(' '):
            prefix += ' '
        self.prefix = prefix

    def _is_internal_frame(self, frame: types.FrameType) -> bool:
        return frame.f_code.co_filename == self.__class__.__aenter__.__code__.co_filename

    @staticmethod
    def _get_source(filename: str, lineno: int) -> str:
        with contextlib.suppress(Exception):
            return linecache.getline(filename, lineno).strip()
        return ""

    def trace(self, frame: types.FrameType, event: str, _: Any):
        """ trace frame and event """
        if frame not in self.target_frames:
            if self.depth == 1 or self._is_internal_frame(frame):
                return None
            else:
                _frame_candidate = frame
                for i in range(1, self.depth):
                    _frame_candidate = _frame_candidate.f_back
                    if _frame_candidate is None:
                        return None
                    elif _frame_candidate in self.target_frames:
                        break
                else:
                    return None

        if not hasattr(_TLocal, 'depth'):
            _TLocal.depth = 0

        if event == 'call':
            _TLocal.depth += 1

        # format output
        if self._handler and self._formatter:
            lineno = frame.f_lineno
            filename = frame.f_code.co_filename
            func_name = frame.f_code.co_name
            source_line = self._get_source(filename, lineno)

            if event == 'line':
                self._handler.write_line(
                    event='line',
                    filename=filename,
                    lineno=lineno,
                    func_name=func_name,
                    source_line=source_line,
                )

                old_local_reprs = self.frame_to_local_reprs.get(frame, {})
                current_locals = {}
                for name, value in frame.f_locals.items():
                    if not name.startswith('_'):
                        current_locals[name] = value

                for name, value in current_locals.items():
                    if name not in old_local_reprs:
                        self._handler.write_var(var_type='var', name=name, value=value)
                    elif old_local_reprs[name] != value:
                        self._handler.write_var(var_type='modified', name=name, value=value)

                self.frame_to_local_reprs[frame] = current_locals.copy()

            elif event == 'call':
                self._handler.increase_indent()
                self._handler.write_line(
                    event='call',
                    filename=filename,
                    lineno=lineno,
                    func_name=func_name,
                )

            elif event == 'return':
                self._handler.write_line(
                    event='return',
                    filename=filename,
                    lineno=lineno,
                    func_name=func_name,
                )
                self._handler.decrease_indent()
                self.frame_to_local_reprs.pop(frame, None)
                _TLocal.depth -= 1

        return self.trace

    async def __aenter__(self):
        if not hasattr(_TLocal, 'depth'):
            _TLocal.depth = -1

        # 获取调用者帧
        calling_frame = inspect.currentframe().f_back

        if not self._is_internal_frame(calling_frame):
            calling_frame.f_trace = self.trace
            self.target_frames.add(calling_frame)

        # save original_trace_functions
        stack = getattr(self.thread_local, 'original_trace_functions', [])
        stack.append(sys.gettrace())
        self.thread_local.original_trace_functions = stack
        self.start_times[calling_frame] = datetime.now()

        # open output
        stream = open_output_stream(self.output_spec, self.overwrite)
        self._stream = stream
        self._opened_by_us = self.output_spec is not None and isinstance(self.output_spec, (str, Path))

        # create formatter and handler
        self._formatter = OutputFormatter(
            prefix=self.prefix,
            include_thread_info=self.thread_info,
            use_colors=self.color,
        )
        writer = OutputWriter(stream)
        self._handler = OutputHandler(writer, self._formatter)

        sys.settrace(self.trace)

        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        # pep original_trace_functions
        stack = getattr(self.thread_local, 'original_trace_functions', [])
        if stack:
            sys.settrace(stack.pop())

        calling_frame = inspect.currentframe().f_back
        self.target_frames.discard(calling_frame)
        self.frame_to_local_reprs.pop(calling_frame, None)
        self.start_times.pop(calling_frame, None)

        # close output
        if hasattr(self, '_stream') and self._stream and self._opened_by_us:
            close_output_stream(self._stream, True)

        return False
