from __future__ import annotations

import pdb
import sys
import os

from textwrap import indent
from functools import wraps

from .agents import FSM, MetaFSM


def wrapcmd(func):
    @wraps(func)
    def wrapper(self, arg: str, temporary=False):
        sys.settrace(self.trace_dispatch)

        known = globals()
        known.update(self.curframe.f_globals)
        known.update(self.curframe.f_locals)
        known["agent"] = known.get("self", None)
        known["model"] = known.get("self", {}).get("model")
        known["attrs"] = arg.strip().split()

        exec(func.__code__, known, known)

    return wrapper


class Debug(pdb.Pdb):
    def __init__(self, *args, skip_soil=False, **kwargs):
        skip = kwargs.get("skip", [])
        if skip_soil:
            skip.append("soil")
            skip.append("contextlib")
            skip.append("soil.*")
            skip.append("mesa.*")
        super(Debug, self).__init__(*args, skip=skip, **kwargs)
        self.prompt = "[soil-pdb] "

    @staticmethod
    def _soil_agents(model, attrs=None, pretty=True, **kwargs):
        for agent in model.agents(**kwargs):
            d = agent
            print(" - " + indent(agent.to_str(keys=attrs, pretty=pretty), "  "))

    @wrapcmd
    def do_soil_agents():
        return Debug._soil_agents(model, attrs=attrs or None)

    do_sa = do_soil_agents

    @wrapcmd
    def do_soil_list():
        return Debug._soil_agents(model, attrs=["state_id"], pretty=False)

    do_sl = do_soil_list

    def do_continue_state(self, arg):
        self.do_break_state(arg, temporary=True)
        return self.do_continue("")

    do_cs = do_continue_state

    @wrapcmd
    def do_soil_agent():
        if not agent:
            print("No agent available")
            return

        keys = None
        if attrs:
            keys = []
            for k in attrs:
                for key in agent.keys():
                    if key.startswith(k):
                        keys.append(key)

        print(agent.to_str(pretty=True, keys=keys))

    do_aa = do_soil_agent

    def do_break_state(self, arg: str, instances=None, temporary=False):
        """
        Break before a specified state is stepped into.
        """

        klass = None
        state = arg
        if not state:
            self.error("Specify at least a state name")
            return

        state, *tokens = state.lstrip().split()
        if tokens:
            instances = list(eval(token) for token in tokens)

        colon = state.find(":")

        if colon > 0:
            klass = state[:colon].rstrip()
            state = state[colon + 1 :].strip()

            print(klass, state, tokens)
            klass = eval(klass, self.curframe.f_globals, self.curframe_locals)

        if klass:
            klasses = [klass]
        else:
            klasses = [
                k
                for k in self.curframe.f_globals.values()
                if isinstance(k, type) and issubclass(k, FSM)
            ]

        if not klasses:
            self.error("No agent classes found")

        for klass in klasses:
            try:
                func = getattr(klass, state)
            except AttributeError:
                self.error(f"State {state} not found in class {klass}")
                continue
            if hasattr(func, "__func__"):
                func = func.__func__

            code = func.__code__
            # use co_name to identify the bkpt (function names
            # could be aliased, but co_name is invariant)
            funcname = code.co_name
            lineno = code.co_firstlineno
            filename = code.co_filename

            # Check for reasonable breakpoint
            line = self.checkline(filename, lineno)
            if not line:
                raise ValueError("no line found")
                # now set the break point
            cond = None
            if instances:
                cond = f"self.unique_id in { repr(instances) }"

            existing = self.get_breaks(filename, line)
            if existing:
                self.message("Breakpoint already exists at %s:%d" % (filename, line))
                continue
            err = self.set_break(filename, line, temporary, cond, funcname)
            if err:
                self.error(err)
            else:
                bp = self.get_breaks(filename, line)[-1]
                self.message("Breakpoint %d at %s:%d" % (bp.number, bp.file, bp.line))

    do_bs = do_break_state

    def do_break_state_self(self, arg: str, temporary=False):
        """
        Break before a specified state is stepped into, for the current agent
        """
        agent = self.curframe.f_locals.get("self")
        if not agent:
            self.error("No current agent.")
            self.error("Try this again when the debugger is stopped inside an agent")
            return

        arg = f"{agent.__class__.__name__}:{ arg } {agent.unique_id}"
        return self.do_break_state(arg)

    do_bss = do_break_state_self


debugger = None


def set_trace(frame=None, **kwargs):
    global debugger
    if debugger is None:
        debugger = Debug(**kwargs)
    frame = frame or sys._getframe().f_back
    debugger.set_trace(frame)


def post_mortem(traceback=None, **kwargs):
    global debugger
    if debugger is None:
        debugger = Debug(**kwargs)
    t = sys.exc_info()[2]
    debugger.reset()
    debugger.interaction(None, t)
