import streamlit as st
import re
import uuid
import subprocess
import sys
from datetime import datetime
from threading import Thread
from queue import Queue
from pygmyclaw import PygmyClaw

st.set_page_config(page_title="Py", layout="wide")

# --------------------------------------------------
# GLOBAL STATE
# --------------------------------------------------

if "job_store" not in st.session_state:
    st.session_state.job_store = {}

if "job_queue" not in st.session_state:
    st.session_state.job_queue = Queue()

if "agent" not in st.session_state:
    st.session_state.agent = PygmyClaw()

if "worker_started" not in st.session_state:
    st.session_state.worker_started = False

job_store = st.session_state.job_store
job_queue = st.session_state.job_queue
agent = st.session_state.agent


# --------------------------------------------------
# BACKGROUND WORKER
# --------------------------------------------------

def worker():
    while True:
        job_id = job_queue.get()
        job = job_store[job_id]
        job["status"] = "running"
        try:
            system_prompt = f"""
You are a Python coding assistant.

Return ONLY runnable Python code.
Do NOT explain anything.

Task:
{job['prompt']}
"""
            result = agent.submit_prompt(system_prompt, job["tool"])

            # ===== FIX =====
            match = re.search(r"```python(.*?)```", result, re.S)
            if match:
                code = match.group(1).strip()
            else:
                code = result.strip()
            # ===============

            job["response"] = result
            job["code"] = code
            job["status"] = "completed"

        except Exception as e:
            job["status"] = "failed"
            job["response"] = str(e)

        job_queue.task_done()

# --------------------------------------------------
# START WORKER
# --------------------------------------------------

if not st.session_state.worker_started:

    Thread(target=worker, daemon=True).start()

    st.session_state.worker_started = True


# --------------------------------------------------
# HEADER
# --------------------------------------------------

st.title("🧠 PygmyClaw AI Dev Dashboard")

# --------------------------------------------------
# CREATE TASK
# --------------------------------------------------

with st.expander("🚀 Create AI Task", expanded=True):

    col1, col2 = st.columns([4,1])

    prompt = col1.text_area(
        "Prompt",
        height=120,
        placeholder="write a python function to add two numbers",
        value="write a python function to add two numbers"
    )

    tool = col2.selectbox(
        "Tool",
        ["AI Agent"]
    )

    if col2.button("Create Job"):

        job_id = str(uuid.uuid4())[:8]

        job_store[job_id] = {
            "prompt": prompt,
            "tool": tool,
            "status": "queued",
            "response": "",
            "code": "",
            "created": datetime.now().strftime("%H:%M:%S")
        }

        job_queue.put(job_id)

        st.success(f"Job {job_id} queued")


# --------------------------------------------------
# METRICS
# --------------------------------------------------

queued = sum(1 for j in job_store.values() if j["status"] == "queued")
running = sum(1 for j in job_store.values() if j["status"] == "running")
done = sum(1 for j in job_store.values() if j["status"] == "completed")
failed = sum(1 for j in job_store.values() if j["status"] == "failed")

col1, col2, col3, col4 = st.columns(4)

col1.metric("Queued", queued)
col2.metric("Running", running)
col3.metric("Completed", done)
col4.metric("Failed", failed)

st.divider()

# --------------------------------------------------
# JOB DASHBOARD
# --------------------------------------------------

if not job_store:
    st.info("No jobs yet")

for job_id, job in list(job_store.items())[::-1]:

    with st.container():

        st.subheader(f"Job {job_id}")

        st.write("Status:", job["status"])
        st.write("Created:", job["created"])

        st.code(job["prompt"], language="text")

        # -----------------------
        # AI RESPONSE
        # -----------------------

        if job.get("response"):

            st.markdown("### 🤖 AI Response")

            st.write(job["response"])

        # -----------------------
        # CODE EDITOR (ALWAYS VISIBLE)
        # -----------------------

        st.markdown("### 💻 Generated Code")

        job["code"] = st.text_area(
            "Edit Code",
            value=job.get("code", ""),
            height=220,
            key=f"code_{job_id}"
        )

        col1, col2 = st.columns(2)

        # -----------------------
        # RUN CODE
        # -----------------------

        if col1.button("▶ Run Code", key=f"run_{job_id}"):

            try:

                code = job.get("code", "")

                # detect imports
                imports = re.findall(
                    r"^\s*(?:import|from)\s+([\w_]+)",
                    code,
                    flags=re.MULTILINE
                )

                # auto install missing packages
                for pkg in imports:
                    try:
                        __import__(pkg)
                    except ImportError:
                        subprocess.run(
                            [sys.executable, "-m", "pip", "install", pkg],
                            check=True
                        )

                local_vars = {}

                exec(code, {}, local_vars)

                st.success("Execution Output")

                st.json(local_vars)

            except Exception as e:

                st.error(str(e))

        # -----------------------
        # DELETE JOB
        # -----------------------

        if col2.button("Delete Job", key=f"del_{job_id}"):

            del job_store[job_id]

            st.rerun()

        st.divider()