from __future__ import annotations
import asyncio
import logging
import os
import threading
import websockets
import queue
from relationalai import Resources, debugging
from relationalai.clients.config import PUBLIC_CONFIG_KEYS, Config # to get the users root file path
import __main__

from relationalai.clients.profile_polling import ProfilePollerHandler

class DummyRecord(logging.LogRecord):
    def __init__(self, msg):
        self.msg = msg

#------------------------------------------------------------------------------
# Logging Handler
#------------------------------------------------------------------------------

class WebSocketLoggingHandler(logging.Handler):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        from relationalai.analysis.mechanistic import Mechanism
        self.Mechanism = Mechanism
        self.queue: queue.Queue[str] = queue.Queue()

    def start(self, url: str, config: Config):
        main_path = getattr(__main__, "__file__", None)

        # Catch any spans that are the programs parent that were added before the debug client was started.
        open_spans = []
        existing_parent = debugging.TIP_SPAN
        while existing_parent:
            open_spans.append(existing_parent)
            existing_parent = existing_parent.parent
        open_spans.reverse()

        program_span = None

        for span in open_spans:
            if span.type == "program":
                program_span = span
            self.queue.put_nowait(self.format_raw({"event": "span_start", "span": span}))

        if not program_span:
            program_span = debugging.span_start(
                "program",
                main=os.path.relpath(main_path) if main_path else "notebook",
                **{k: config.get(k, None) for k in PUBLIC_CONFIG_KEYS if config.get(k, None) is not None}
            )

        self.thread = threading.Thread(target=self.connect, args=(url, program_span))
        self.thread.start()

    def emit(self, record):
        d = record.msg
        if (isinstance(d, dict) and
            d["event"] == "span_start" and
            "task" in d["span"].attrs and
            "mech" not in d["span"].attrs):
            d["span"].attrs["mech"] = self.Mechanism(d["span"].attrs["task"])
        log_entry = self.format(record)
        self.queue.put_nowait(log_entry)

    def format_raw(self, msg):
        return self.format(DummyRecord(msg))

    def connect(self, url: str, span: debugging.Span):
        try:
            try:
                loop = asyncio.get_running_loop()
                loop.run_until_complete(self.connect_async(url, span))
            except RuntimeError:
                asyncio.run(self.connect_async(url, span))
        except (ConnectionRefusedError, websockets.WebSocketException, OSError, asyncio.TimeoutError):
            pass

    async def connect_async(self, url: str, span: debugging.Span):
        async with websockets.connect(url) as ws:
            print(f"Connected to debugger at {url}...")
            alive = True
            while True:
                if not threading.main_thread().is_alive():
                    alive = False
                    debugging.span_flush()

                try:
                    log_entry = self.queue.get(timeout=1)
                    await ws.send(log_entry)
                except queue.Empty:
                    if not alive:
                        break
                except websockets.ConnectionClosedError:
                    break

#------------------------------------------------------------------------------
# Enable debugging
#------------------------------------------------------------------------------

already_debugging = False
def start_debugger_session(config: Config, host: str|None = None, port: int|None = None):
    global already_debugging
    if already_debugging:
        return

    if host is None:
        host = "0.0.0.0"

    if port is None:
        port = 8080

    ws_handler = WebSocketLoggingHandler()
    ws_handler.setFormatter(debugging.JsonFormatter())
    debugging.logger.addHandler(ws_handler)
    ws_handler.start(f"ws://{host}:{port}/ws/program", config)
    
    resources = Resources(config=config)
    debugging.logger.addHandler(ProfilePollerHandler(resources))

    already_debugging = True
