"""Code execution engine — runs user code in a sandbox with live previews."""
import sys
import threading
import traceback
import queue as _queue

import gradio as gr
import torch

from models import _cached, _load_text_model, _load_image_model, _load_tts_model

# Global state for cancellation
_current_input_queue = None
_active_run = None  # stores {'q': queue, 'input_q': queue, 'thread': thread, 'cancel_event': event}


def run_custom_code(code):
    global _current_input_queue, _active_run
    import io as _io, sys as _sys, subprocess as _sp, traceback as _tb, \
           threading as _threading, queue as _queue, os as _os

    cancel_event = threading.Event()
    if _active_run:
        _active_run['cancel_event'].set()
        try:
            _active_run['input_q'].put("__CANCELLED__")
        except Exception:
            pass
    q = _queue.Queue()
    _input_queue = _queue.Queue()
    _current_input_queue = _input_queue

    _IMG = {".png",".jpg",".jpeg",".gif",".bmp",".webp",".svg"}
    _AUD = {".mp3",".wav",".ogg",".flac",".m4a",".aac",".wma"}
    _VID = {".mp4",".avi",".mov",".mkv",".webm",".m4v"}

    def _preview(data):
        if isinstance(data, list) and data and isinstance(data[0], dict):
            md = "\n\n".join(
                f"**{m.get('role','user').title()}:** {m.get('content','')}"
                for m in data
            )
            q.put(("prev", "txt", md))
            q.put(("input_wait", None))
            # Check cancel while waiting for user input
            while True:
                if cancel_event.is_set():
                    return "__CANCELLED__"
                try:
                    return _input_queue.get(timeout=0.5)
                except _queue.Empty:
                    continue
        if not isinstance(data, str):
            q.put(("prev", "txt_append", str(data)))
            return
        ext = _os.path.splitext(data)[1].lower()
        if _os.path.isfile(data):
            if ext in _IMG:
                q.put(("prev", "img", data)); return
            if ext in _AUD:
                q.put(("prev", "aud", data)); return
            if ext in _VID:
                q.put(("prev", "vid", data)); return
            try:
                with open(data, "r") as f:
                    q.put(("prev", "txt", f.read()))
            except Exception:
                q.put(("out", "[preview error]\n"))
            return
        q.put(("prev", "txt_append", data))

    def _preview_text(content, clear=False):
        q.put(("prev", "txt" if clear else "txt_append", content))

    env = {
        "torch": torch,
        "preview": _preview,
        "preview_text": _preview_text,
        "load_text_model": _cached(_load_text_model),
        "load_image_model": _cached(_load_image_model),
        "load_tts_model": _cached(_load_tts_model),
        "_cancel_event": cancel_event,
        "__builtins__": __builtins__,
    }

    class _StdoutIO:
        def write(self, s):
            q.put(("out", s))
        def flush(self):
            pass

    class _StderrIO:
        def write(self, s):
            q.put(("err", s))
        def flush(self):
            pass

    old_out, old_err = _sys.stdout, _sys.stderr
    _sys.stdout = _StdoutIO()
    _sys.stderr = _StderrIO()

    def _exec_thread():
        try:
            py_blocks = []
            for line in code.split("\n"):
                if cancel_event.is_set():
                    return
                if line.startswith("#!"):
                    if py_blocks:
                        exec("\n".join(py_blocks), env)
                        py_blocks = []
                    cmd = line[2:].strip()
                    result = _sp.run(cmd, shell=True, capture_output=True, text=True)
                    if result.stdout:
                        print(result.stdout, end="")
                    if result.stderr:
                        print(result.stderr, file=_sys.stderr, end="")
                else:
                    py_blocks.append(line)
            if py_blocks and not cancel_event.is_set():
                exec("\n".join(py_blocks), env)
        except Exception:
            if not cancel_event.is_set():
                print(_tb.format_exc(), file=_sys.stderr)
        finally:
            if not cancel_event.is_set():
                q.put(None)

    thread = _threading.Thread(target=_exec_thread, daemon=True)
    _active_run = {'q': q, 'input_q': _input_queue, 'thread': thread, 'cancel_event': cancel_event}
    thread.start()

    out_acc, err_acc = "", ""
    p_vals = {"img": None, "aud": None, "vid": None, "txt": None}
    p_vis  = {"img": False, "aud": False, "vid": False, "txt": False}

    def _esc(s):
        return s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

    def _build_console():
        inner = ""
        if out_acc:
            inner += f"<pre>{_esc(out_acc)}</pre>"
        if err_acc:
            inner += f"<pre style='font-style:italic;color:#e74c3c'>{_esc(err_acc)}</pre>"
        if not inner:
            return ""
        return f"<div class='cscroll' style='max-height:300px;overflow-y:auto;font-family:monospace;font-size:0.85em;padding:8px;border-radius:6px;'>{inner}</div>"

    def _yield(running=True):
        return (
            _build_console(),
            gr.update(visible=p_vis["img"], value=p_vals["img"]) if p_vis["img"] else gr.update(visible=False),
            gr.update(visible=p_vis["aud"], value=p_vals["aud"]) if p_vis["aud"] else gr.update(visible=False),
            gr.update(visible=p_vis["vid"], value=p_vals["vid"]) if p_vis["vid"] else gr.update(visible=False),
            gr.update(visible=p_vis["txt"], value=p_vals["txt"]) if p_vis["txt"] else gr.update(visible=False),
            gr.update(visible=chat_visible),
            gr.update(visible=not running),  # Run button
            gr.update(visible=running),  # Stop button
        )

    chat_visible = False

    yield _yield(running=True)

    def _show(kind, val):
        for k in p_vis:
            p_vis[k] = (k == kind)
        if kind == "txt":
            p_vals[kind] = str(val)
        elif kind == "txt_append":
            p_vis["txt"] = True
            p_vals["txt"] = (p_vals["txt"] or "") + str(val)
        else:
            p_vals[kind] = val

    try:
        while True:
            try:
                msg = q.get(timeout=0.5)
            except _queue.Empty:
                if cancel_event.is_set():
                    break
                continue
            if msg is None:
                break
            if cancel_event.is_set():
                break
            t = msg[0]
            if t == "out":
                out_acc += msg[1]
                chat_visible = False
                yield _yield()
            elif t == "err":
                err_acc += msg[1]
                chat_visible = False
                yield _yield()
            elif t == "input_wait":
                chat_visible = True
                yield _yield()
            elif t == "prev":
                _show(msg[1], msg[2])
                chat_visible = False
                yield _yield()
    except Exception:
        _tb.print_exc()
    finally:
        cancel_event.set()
        _current_input_queue = None
        _sys.stdout, _sys.stderr = old_out, old_err
        _active_run = None
        chat_visible = False
        yield _yield(running=False)


def cancel_run():
    """Called when Stop button is clicked - cancels current run, returns UI updates."""
    global _active_run
    if _active_run:
        _active_run['cancel_event'].set()
        # Unblock input queue if chat is waiting
        if _active_run.get('input_q'):
            try:
                _active_run['input_q'].put("__CANCELLED__")
            except Exception:
                pass
    return (
        gr.update(visible=True),   # Run button visible
        gr.update(visible=False),  # Stop button hidden
    )