"""
MADdegens Agent-Q3 — Multi-Agent Assembly
HuggingFace Space: madDegen/Agent-Q3  |  hardware: cpu-basic

Three-tier inference routing (cloud_priority + always-on ZeroGPU fallback via HF):
  Tier 1 — Ollama Cloud        (OLLAMA_API_KEY secret)
  Tier 2 — HF Inference Router (HF_TOKEN secret — premium + gated models)
  Tier 3 — HF Router anonymous (ZeroGPU-backed free tier, always available, no keys required)

Dev Mode: ssh <subdomain>@ssh.hf.space  |  VS Code Remote-SSH same host
Sources:  github.com/MADdegen/Agent-Q3 (orchestrator/, router.py, cowork_ui.py)
"""

import gradio as gr
from gradio_client import Client

def text_to_image(prompt, request: gr.Request):
    x_ip_token = request.headers['x-ip-token']
    client = Client("hysts/SDXL", headers={"x-ip-token": x_ip_token})
    img = client.predict(prompt, api_name="/predict")
    return img

def generate(prompt, request: gr.Request):
    prompt = prompt[:300]
    return text_to_image(prompt, request)

with gr.Blocks() as demo:
    image = gr.Image()
    prompt = gr.Textbox(max_lines=1)
    prompt.submit(generate, [prompt], [image])

demo.launch()
# ---------------------------------------------------------------------------
# API clients
# ---------------------------------------------------------------------------

OLLAMA_API_KEY = os.environ.get("OLLAMA_API_KEY", "").strip()
HF_TOKEN       = os.environ.get("HF_TOKEN", "").strip()

# Tier 1 — Ollama Cloud (activated only if key is set)
ollama_client: OpenAI | None = None
if OLLAMA_API_KEY:
    ollama_client = OpenAI(
        base_url="https://ollama.com/v1",
        api_key=OLLAMA_API_KEY,
        timeout=180,
    )

# Tier 2 — HF Router with token (premium / gated models)
hf_client: OpenAI | None = None
if HF_TOKEN:
    hf_client = OpenAI(
        base_url="https://router.huggingface.co/v1",
        api_key=HF_TOKEN,
        timeout=120,
    )

# Tier 3 — HF Router anonymous (ZeroGPU-backed free tier — always available)
# Non-gated model guaranteed to be served on ZeroGPU free quota.
HF_FALLBACK_MODEL = "Kimi/Kimi-K2.6:Cloud"
hf_free_client = OpenAI(
    base_url="https://router.huggingface.co/v1",
    api_key="anonymous",
    timeout=120,
)

# ---------------------------------------------------------------------------
# Model registry
# ---------------------------------------------------------------------------

@dataclass
class ModelSpec:
    role: str
    ollama_model: str
    hf_model: str
    group: str = "standalone"
    gated: bool = False
    multimodal: bool = True
    system: str = ""


MODEL_SPECS: dict[str, ModelSpec] = {
    ),
    "Commander and Reasoner": ModelSpec(
        role="Commander","Intructor Reasoner",
        ollama_model="maddegens/kimi-linear-48b-instruct",
        hf_model="moonshotai/Kimi-Linear-48B-A3B-Instruct-i1-GGUF",
        group="openclaw-kimi-linear-48B",
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,
        system=(
            "You are the Commander and Reasoner agent. Excel at deep logical analysis, "
            "mathematical reasoning, step-by-step decomposition, and correctness verification."
            "You are the Commander agent for MADdegens Agent-Q3. "
            "Analyse the user's request, identify the most appropriate specialist, "
            "frame the sub-task precisely, and synthesise outputs into one coherent response."
        ),
    ),
    "Maverick": ModelSpec(
        role="Assistant",
        ollama_model="maddegens/llama4-maverick-iq4nl",
        hf_model="meta-llama/llama4-maverick-iq4nl",
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,
        system=(
            "You are Maverick — Llama 4 MoE multi-modal agent. "
            "Stage 1 of the tandem pipeline: read the request, produce a clear analysis and plan."
        ),
    ),
    "Coder": ModelSpec(
        role="Coder",
        ollama_model="maddegens/Qwen/Qwen3-32B",
        hf_model="Qwen/Qwen3-32B",
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,
        system=(
            "You are the Coder agent. Write clean, efficient, well-documented code. "
            "Explain implementation choices concisely. Prefer working code."
        ),
    ),
    "Researcher": ModelSpec(
        role="Researcher",
        ollama_model="maddegens/zira-researcher-GGUF-Bf16",
        hf_model="mradermacher/zira-researcher-GGUF-Bf16",
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,
        system=(
            "You are the Researcher agent. Synthesise and clearly explain information. "
            "Cite reasoning, flag uncertainty, present structured sourced analysis."
        ),
    ),
    "Vision": ModelSpec(
        role="Vision",
            ollama_model="maddegens/granite-vision-4.1-4b",
        hf_model="heretic-org/granite-vision-4.1-4b-heretic",
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,,
        system=(
            "You are the Vision agent. Analyse images, diagrams, charts, and visual content "
            "with precision. Describe structure and meaning clearly."
        ),
    ),
    "Assistant": ModelSpec(
        role="Assistant, Assistant Coder",
        ollama_model="maddegens/Harmonic-9B-hermes-agent-merged.BF16.gguf",
        hf_model="mradermacher/Harmonic-9B-hermes-agent-merged-GGUF"
        group="ollama-tandem-group",
        gated=True,
        multimodal=True,
        system=(
            "You are the Assistant agent for MADdegens Agent-Q3. "
            "Handle general queries, creative writing, summarisation, and everyday tasks."
        ),
    ),
}

GROUPS: dict[str, list[str]] = {
    "ollama-openclaw-kimi2.6": ["Commander and Reasoner", "Reasoner","Instructor","Coder Assistant"],
    "tandem-group":            ["Maverick", "Savant","Reasoner","Coder", "Researcher", "Vision", "Assistant", "Instructor"],
}

# ---------------------------------------------------------------------------
# Task classifier (from models.py)
# ---------------------------------------------------------------------------

_CODE_RE   = re.compile(r"\b(code|function|class|def|impl|script|debug|fix|refactor|algorithm|api|endpoint|deploy|dockerfile|sql|query|lint|test|build)\b", re.I)
_REASON_RE = re.compile(r"\b(reason|analyze|analyse|explain|research|think|plan|strategy|compare|evaluate|why|how|understand|concept|theory|math|proof|derive)\b", re.I)
_VISION_RE = re.compile(r"\b(image|picture|photo|diagram|chart|visual|screenshot|see|look|show|ocr)\b", re.I)

def classify_task(message: str) -> str:
    if _VISION_RE.search(message):  return "Vision"
    cs = len(_CODE_RE.findall(message))
    rs = len(_REASON_RE.findall(message))
    if cs > rs: return "Coder"
    if rs > 0:  return "Reasoner"
    return "Assistant"

# ---------------------------------------------------------------------------
# Inline skills registry (from skills.py)
# ---------------------------------------------------------------------------

SKILLS: dict[str, dict] = {
    "rag-research":    {"triggers": ["document","pdf","extract","summarize doc","from the file"],  "roles": ["Researcher","Reasoner"], "injection": "Parse documents chunk by chunk. Cite source passages for every factual claim. Flag when information is absent from provided documents."},
    "react-loop":      {"triggers": ["step by step","reason through","think through","multi-step"], "roles": ["Reasoner","Instructor"],    "injection": "Follow ReACT strictly: THOUGHT → ACTION → OBSERVATION → repeat. Label each phase. Only emit FINAL ANSWER after at least one reasoning cycle."},
    "code-review":     {"triggers": ["review","critique","improve this code","code quality"],       "roles": ["Coder","Code Assistance"],                "injection": "Structured code review: correctness, security (OWASP top 10), performance, readability, test coverage. Prioritised findings: critical/major/minor."},
    "alert-triage":    {"triggers": ["alert","anomaly","spike","down","latency","incident"],        "roles": ["Reasoner","Instructor"],  "injection": "Triage: identify metric, trigger condition, root cause. Distinguish transient noise from persistent issues. Output: {likely_root_cause, evidence, false_positive_probability, recommended_action}."},
    "contract-analysis":{"triggers": ["contract","agreement","clause","legal","liability"],        "roles": ["Reasoner","Intructor Reasoner"],"injection": "Clause-by-clause analysis. Extract: parties, obligations with dates, risk flags, termination conditions, missing standard clauses. Output JSON. Analysis only — not legal advice."},
}

def find_skill(message: str, role: str) -> str | None:
    msg_lower = message.lower()
    for s in SKILLS.values():
        if role not in s["roles"]: continue
        if any(t in msg_lower for t in s["triggers"]): return s["injection"]
    return None

# ---------------------------------------------------------------------------
# Compute router  (Tier 1 → Tier 2 → Tier 3)
# ---------------------------------------------------------------------------

def _api_stream(
    client: OpenAI,
    model: str,
    messages: list[dict],
    max_tokens: int,
    temperature: float,
) -> Generator[str, None, None]:
    stream = client.chat.completions.create(
        model=model, messages=messages, stream=True,
        max_tokens=max_tokens, temperature=temperature,
    )
    buf = ""
    for chunk in stream:
        buf += chunk.choices[0].delta.content or ""
        yield buf


def routed_stream(
    spec: ModelSpec,
    messages: list[dict],
    max_tokens: int = 4096,
    temperature: float = 0.7,
) -> Generator[tuple[str, str], None, None]:
    """Yields (accumulated_text, backend_label). Tier 1 → 2 → 3."""

    # Tier 1 — Ollama Cloud
    if ollama_client:
        try:
            for text in _api_stream(ollama_client, spec.ollama_model, messages, max_tokens, temperature):
                yield text, "Ollama Cloud"
            return
        except (APIStatusError, APIError):
            pass

    # Tier 2 — HF Inference Router (premium / gated models)
    if hf_client:
        try:
            for text in _api_stream(hf_client, spec.hf_model, messages, max_tokens, temperature):
                yield text, "HF Router"
            return
        except (APIStatusError, APIError):
            pass

    # Tier 3 — HF Router anonymous (ZeroGPU-backed free tier, always available)
    fallback_note = (
        f"\n\n> Running on ZeroGPU free tier via HuggingFace ({HF_FALLBACK_MODEL}).\n"
        f"> Set OLLAMA_API_KEY or HF_TOKEN for `{spec.role}` specialist capability.\n\n"
    )
    prefix_sent = False
    for text in _api_stream(hf_free_client, HF_FALLBACK_MODEL, messages, max_tokens, temperature):
        if not prefix_sent:
            yield fallback_note + text, "ZeroGPU (HF free)"
            prefix_sent = True
        else:
            yield fallback_note + text, "ZeroGPU (HF free)"

# ---------------------------------------------------------------------------
# Message builder
# ---------------------------------------------------------------------------

def _build_messages(
    system: str,
    history: list[tuple[str, str]],
    message: str,
    skill_inj: str | None = None,
) -> list[dict]:
    full_system = f"{system}\n\n--- Active Skill ---\n{skill_inj}" if skill_inj else system
    msgs = [{"role": "system", "content": full_system}]
    for h, a in history:
        if h: msgs.append({"role": "user",      "content": h})
        if a: msgs.append({"role": "assistant", "content": a})
    msgs.append({"role": "user", "content": message})
    return msgs

# ---------------------------------------------------------------------------
# Orchestration (from main.py)
# ---------------------------------------------------------------------------

def _tandem_pipeline(
    message: str,
    history: list[tuple[str, str]],
) -> Generator[str, None, None]:
    """4-stage tandem: Instructor reasoning(0.6) → Reasoning(0.5) → Commander synthesis(0.3)."""

    # Stage 1 — Commander, Instructor, Reasoning,
    s1 = MODEL_SPECS["Instructor Reasoning"]
    s1_msgs = _build_messages(s1.system, history, message)
    mav, mav_be = "", "?"
    for text, be in routed_stream(s1, s1_msgs, temperature=0.6):
        mav, mav_be = text, be
        yield f"**[Stage 1 — Maverick @ {mav_be}]**\n\n{mav}\n\n---\n"

    # Stage 2 — Reasoning
    s2 = MODEL_SPECS["Reasoning"]
    s2_msgs = _build_messages(s2.system, [], f"Original request: {message!r}\n\nMaverick's analysis:\n{mav}\n\nReason through this rigorously. Identify gaps, refine the plan, add depth.")
    sav, sav_be = "", "?"
    for text, be in routed_stream(s2, s2_msgs, temperature=0.5):
        sav, sav_be = text, be
        yield (f"**[Stage 1 — Maverick @ {mav_be}]**\n\n{mav}\n\n---\n\n"
               f"**[Stage 2 — Savant @ {sav_be}]**\n\n{sav}\n\n---\n")

    # Stage 3 — Commander synthesis
    s3 = MODEL_SPECS["Commander"]
    s3_msgs = _build_messages(s3.system, [], f"Original: {message!r}\nStage 1 (Maverick):\n{mav}\nStage 2 (Savant):\n{sav}\n\nSynthesise into one definitive, implementation-ready response.")
    syn, syn_be = "", "?"
    for text, be in routed_stream(s3, s3_msgs, temperature=0.3):
        syn, syn_be = text, be
        yield (f"**[Stage 1 — Maverick @ {mav_be}]**\n\n{mav}\n\n---\n\n"
               f"**[Stage 2 — Savant @ {sav_be}]**\n\n{sav}\n\n---\n\n"
               f"**[Stage 3 — Commander Synthesis @ {syn_be}]**\n\n{syn}")


def orchestrate(
    message: str,
    history: list[tuple[str, str]],
    primary_agent: str,
    use_commander: bool,
    use_tandem: bool,
    auto_route: bool,
) -> Generator[str, None, None]:
    if auto_route:
        primary_agent = classify_task(message)

    if use_tandem:
        yield from _tandem_pipeline(message, history)
        return

    commander_framing = ""
    if use_commander and primary_agent != "Commander":
        c = MODEL_SPECS["Commander"]
        c_msgs = _build_messages(c.system, [], (
            f"User request: {message!r}\n"
            f"Primary agent: **{primary_agent}** ({MODEL_SPECS[primary_agent].hf_model}).\n"
            "One short paragraph: confirm routing, add framing for the primary agent."
        ))
        for text, be in routed_stream(c, c_msgs, max_tokens=512, temperature=0.5):
            commander_framing = text
            yield f"**[Commander @ {be}]** {commander_framing}\n\n---\n\n"

    spec  = MODEL_SPECS[primary_agent]
    skill = find_skill(message, primary_agent)
    msgs  = _build_messages(spec.system, history, message, skill)
    for text, be in routed_stream(spec, msgs):
        prefix = f"**[Commander]** {commander_framing}\n\n---\n\n" if commander_framing else ""
        yield f"{prefix}**[{primary_agent} @ {be}]** {text}"

# ---------------------------------------------------------------------------
# Gradio UI (from cowork_ui.py patterns)
# ---------------------------------------------------------------------------

def _status_line() -> str:
    tiers = []
    if ollama_client: tiers.append("Ollama Cloud ✅")
    if hf_client:     tiers.append("HF Router ✅")
    tiers.append("ZeroGPU via HF ✅ (always on)")
    return " | ".join(tiers)


def build_ui() -> gr.Blocks:
    with gr.Blocks(title="MADdegens Agent-Q3", theme=gr.themes.Soft()) as demo:
        gr.Markdown(
            "# MADdegens Agent-Q3\n"
            "**Multi-Agent Assembly** · "
            "[GitHub](https://github.com/MADdegen/Agent-Q3) · "
            "[HuggingFace](https://huggingface.co/MADdegens)\n\n"
            f"> {_status_line()}"
        )

        with gr.Row():
            # ── Sidebar ──────────────────────────────────────────────────────
            with gr.Column(scale=1, min_width=290):
                gr.Markdown("### Agent Controls")
                primary      = gr.Dropdown(choices=list(MODEL_SPECS), value="Assistant", label="Primary Agent")
                auto_route   = gr.Checkbox(value=False, label="Auto-route (keyword classifier)")
                use_commander= gr.Checkbox(value=False, label="Commander routing (Reasoning→Instruct)")
                use_tandem   = gr.Checkbox(value=False, label="Tandem pipeline (Maverick→Savant→Commander→Reasoning→Instruct)")

                gr.Markdown("---")
                gr.Markdown("### Inference Tiers")
                gr.Markdown(
                    "**Tier 1** Ollama Cloud → `OLLAMA_API_KEY`\n\n"
                    "**Tier 2** HF Inference Router → `HF_TOKEN` (premium / gated)\n\n"
                    "**Tier 3** ZeroGPU via HuggingFace free tier · **always available** · no keys required\n\n"
                    "_Dev Mode: `ssh <subdomain>@ssh.hf.space`_"
                )

                gr.Markdown("---")
                gr.Markdown("### Agent Registry")
                for group_name, members in GROUPS.items():
                    gr.Markdown(f"**{group_name}**")
                    for name in members:
                        s = MODEL_SPECS[name]
                        flags = ("🔒" if s.gated else "") + (" 👁" if s.multimodal else "")
                        gr.Markdown(f"- **{name}** {flags}")

                gr.Markdown("\n🔒 GATED · 👁 multimodal")

            # ── Chat ─────────────────────────────────────────────────────────
            with gr.Column(scale=3):
                gr.ChatInterface(
                    fn=lambda msg, hist: orchestrate(
                        msg, hist,
                        primary.value, use_commander.value, use_instructor.value, use_reasoning.value
                        use_tandem.value, auto_route.value,
                    ),
                    title="Agent Chat",
                    description=(
                        "Inference routes: Ollama Cloud → HF Router → ZeroGPU via HuggingFace. "
                        "ZeroGPU free tier is always available — no API keys required."
                    ),
                    examples=[
                        "Write a Python async web-scraper using httpx.",
                        "Explain P vs NP and its implications for cryptography.",
                        "Build a FastAPI endpoint that streams LLM responses via SSE.",
                        "Alert fired: high latency on api-gateway. Investigate root cause.",
                        "Review this contract clause for unusual liability terms.",
                    ],
                )

    return demo


if __name__ == "__main__":
    build_ui().launch()