import os
import json
from dataclasses import dataclass, asdict, field
from typing import Optional, List

import gradio as gr
from openai import OpenAI

# ============================================================
# CONFIG
# ============================================================

OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

if not OPENROUTER_API_KEY:
    raise RuntimeError(
        "OPENROUTER_API_KEY environment variable not found."
    )

MODEL = os.getenv(
    "OPENROUTER_MODEL",
    "openai/gpt-oss-120b:free"
)

client = OpenAI(
    base_url="https://openrouter.ai/api/v1",
    api_key=OPENROUTER_API_KEY
)

# ============================================================
# STATE
# ============================================================

@dataclass
class AgentState:
    name: Optional[str] = None
    role: Optional[str] = None
    goal: Optional[str] = None
    greeting: Optional[str] = None
    tone: Optional[str] = None
    memory: Optional[str] = None
    audience: Optional[str] = None
    domain: Optional[str] = None

    tools: List[str] = field(default_factory=list)
    policies: List[str] = field(default_factory=list)

    def to_dict(self):
        return asdict(self)

# ============================================================
# PROMPTS
# ============================================================

DISCOVERY_PROMPT = """
You are an Agent Architect.

Your job is to help users create AI agents entirely through conversation.

Rules:

- Talk naturally.
- Gather requirements progressively.
- Ask only ONE important follow-up question.
- Never mention prompts.
- Never mention JSON.
- Never mention implementation.
- Focus on understanding the user's vision.

The user's natural language is the source of truth.
"""

ARCHITECT_PROMPT = """
Extract structured information from the user's message.

Return ONLY JSON.

Schema:

{
  "name": null,
  "role": null,
  "goal": null,
  "greeting": null,
  "tone": null,
  "memory": null,
  "audience": null,
  "domain": null,
  "tools": [],
  "policies": []
}
"""

CAPABILITY_PROMPT = """
You are a Capability Planner.

Given an agent specification,
suggest useful capabilities.

Return ONLY JSON.

{
  "tools": []
}
"""

COMPILER_PROMPT = """
You generate system prompts.

Convert the specification into
a production-grade runtime prompt.

Output ONLY the prompt.
"""

# ============================================================
# ORCHESTRATOR
# ============================================================

class AgentOrchestrator:

    def __init__(self):
        self.state = AgentState()
        self.runtime_prompt = ""

    # ========================================================

    def llm(
        self,
        messages,
        max_tokens=500,
        response_format=None
    ):

        kwargs = {
            "model": MODEL,
            "messages": messages,
            "max_tokens": max_tokens
        }

        if response_format:
            kwargs["response_format"] = response_format

        response = client.chat.completions.create(**kwargs)

        content = response.choices[0].message.content

        if not content:
            return ""

        return content.strip()

    # ========================================================

    def update_architecture(
        self,
        user_message: str
    ):

        try:

            result = self.llm(
                [
                    {
                        "role": "system",
                        "content": ARCHITECT_PROMPT
                    },
                    {
                        "role": "user",
                        "content": user_message
                    }
                ],
                response_format={
                    "type": "json_object"
                }
            )

            data = json.loads(result)

            for key, value in data.items():

                if value in [None, "", []]:
                    continue

                if hasattr(self.state, key):
                    setattr(
                        self.state,
                        key,
                        value
                    )

        except Exception as e:
            print("Architect error:", e)

    # ========================================================

    def update_capabilities(self):

        try:

            result = self.llm(
                [
                    {
                        "role": "system",
                        "content": CAPABILITY_PROMPT
                    },
                    {
                        "role": "user",
                        "content": json.dumps(
                            self.state.to_dict(),
                            indent=2
                        )
                    }
                ],
                response_format={
                    "type": "json_object"
                }
            )

            data = json.loads(result)

            tools = data.get("tools", [])

            if isinstance(tools, list):

                merged = set(
                    self.state.tools
                )

                merged.update(tools)

                self.state.tools = list(
                    merged
                )

        except Exception as e:
            print("Capability error:", e)

    # ========================================================

    def compile_runtime_prompt(self):

        try:

            self.runtime_prompt = self.llm(
                [
                    {
                        "role": "system",
                        "content": COMPILER_PROMPT
                    },
                    {
                        "role": "user",
                        "content": json.dumps(
                            self.state.to_dict(),
                            indent=2
                        )
                    }
                ],
                max_tokens=700
            )

        except Exception as e:

            print("Compile error:", e)

    # ========================================================

    def developer_chat(
        self,
        message: str
    ):

        self.update_architecture(message)

        self.update_capabilities()

        self.compile_runtime_prompt()

        reply = self.llm(
            [
                {
                    "role": "system",
                    "content": DISCOVERY_PROMPT
                },
                {
                    "role": "user",
                    "content":
                    f"""
Agent State:

{json.dumps(self.state.to_dict(), indent=2)}

Developer Message:

{message}
"""
                }
            ],
            max_tokens=250
        )

        return reply

    # ========================================================

    def client_chat(
        self,
        message,
        history
    ):

        if not self.runtime_prompt:
            self.compile_runtime_prompt()

        messages = [
            {
                "role": "system",
                "content": self.runtime_prompt
            }
        ]

        for item in history:

            messages.append(
                {
                    "role": item["role"],
                    "content": item["content"]
                }
            )

        messages.append(
            {
                "role": "user",
                "content": message
            }
        )

        return self.llm(
            messages,
            max_tokens=700
        )

# ============================================================
# APP STATE
# ============================================================

orchestrator = AgentOrchestrator()

# ============================================================
# UI CALLBACK
# ============================================================

def chat_handler(
    message,
    history,
    mode
):

    history = history or []

    if mode == "Developer":

        reply = orchestrator.developer_chat(
            message
        )

    else:

        reply = orchestrator.client_chat(
            message,
            history
        )

    history.append(
        {
            "role": "user",
            "content": message
        }
    )

    history.append(
        {
            "role": "assistant",
            "content": reply
        }
    )

    state = orchestrator.state.to_dict()

    return (
        "",
        history,
        state.get("name") or "",
        state.get("role") or "",
        state.get("goal") or "",
        state.get("greeting") or "",
        state.get("memory") or "",
        state.get("tone") or "",
        ", ".join(state.get("tools", [])),
        orchestrator.runtime_prompt
    )

# ============================================================
# UI
# ============================================================

with gr.Blocks(
    title="Agent Platform",
    fill_height=True
) as demo:

    gr.Markdown("# Agent Platform")

    with gr.Row():

        # LEFT PANEL
        with gr.Column(scale=2):

            mode = gr.Dropdown(
                choices=[
                    "Developer",
                    "Client"
                ],
                value="Developer",
                label="Mode"
            )

            chatbot = gr.Chatbot(
                type="messages",
                height=700
            )

            message = gr.Textbox(
                placeholder="Describe your agent..."
            )

        # RIGHT PANEL
        with gr.Column(scale=1):

            gr.Markdown("## Agent Definition")

            name = gr.Textbox(
                label="Name"
            )

            role = gr.Textbox(
                label="Role"
            )

            goal = gr.Textbox(
                label="Goal"
            )

            greeting = gr.Textbox(
                label="Greeting"
            )

            memory = gr.Textbox(
                label="Memory"
            )

            tone = gr.Textbox(
                label="Tone"
            )

            tools = gr.Textbox(
                label="Tools"
            )

            runtime_prompt = gr.Textbox(
                label="Compiled Runtime Prompt",
                lines=18
            )

    message.submit(
        chat_handler,
        inputs=[
            message,
            chatbot,
            mode
        ],
        outputs=[
            message,
            chatbot,
            name,
            role,
            goal,
            greeting,
            memory,
            tone,
            tools,
            runtime_prompt
        ]
    )

# ============================================================
# RUN
# ============================================================

if __name__ == "__main__":
    demo.launch(
        share=True,
        debug=True
    )