"""Lightweight triage using Gemini Flash Lite to decide which model handles a turn."""

from __future__ import annotations

import asyncio
import json
import logging
from enum import Enum
from typing import Any

from google import genai

from epochler.config import get_settings

logger = logging.getLogger(__name__)


class Tier(str, Enum):
    LIGHT = "light"
    HEAVY = "heavy"


_TRIAGE_PROMPT = """\
You are a routing classifier for an autonomous ML research assistant.

Given the latest user message and a short conversation summary, decide how \
much reasoning the response requires.

Reply with EXACTLY one JSON object, nothing else:
{"tier": "light"} or {"tier": "heavy"}

Rules:
- "heavy" if the message asks to write, generate, review, or debug CODE.
- "heavy" if it asks for deep architectural reasoning, novel model design, \
  complex math, or multi-step technical analysis.
- "heavy" if it asks to START a run, REVISE a complex spec section, or \
  design a training/evaluation strategy from scratch.
- "light" for everything else: simple confirmations, approvals, short \
  factual questions, clarifications, parameter tweaks, yes/no answers, \
  greetings, or conversational replies.

When in doubt, pick "heavy" — it is safer to over-think than to under-think.
"""


def _build_triage_input(
    latest_user_message: str,
    recent_summary: str,
) -> str:
    return (
        f"Conversation context (last few turns):\n{recent_summary}\n\n"
        f"Latest user message:\n{latest_user_message}"
    )


def _sync_triage(content: str) -> str:
    settings = get_settings()
    client = genai.Client(api_key=settings.gemini_api_key)
    response = client.models.generate_content(
        model=settings.gemini_lite_model,
        contents=f"{_TRIAGE_PROMPT}\n\n---\n\n{content}",
        config=genai.types.GenerateContentConfig(temperature=0.0),
    )
    return response.text.strip()


async def classify(
    latest_user_message: str,
    recent_messages: list[dict[str, Any]],
) -> Tier:
    """Return the tier that should handle this turn."""
    settings = get_settings()
    if not settings.gemini_api_key:
        return Tier.HEAVY

    summary_parts: list[str] = []
    for msg in recent_messages[-6:]:
        role = msg.get("role", "?")
        text = str(msg.get("content", ""))[:200]
        summary_parts.append(f"{role}: {text}")
    summary = "\n".join(summary_parts) if summary_parts else "(no prior context)"

    content = _build_triage_input(latest_user_message, summary)

    try:
        raw = await asyncio.to_thread(_sync_triage, content)
    except Exception:
        logger.exception("Router triage failed; defaulting to heavy")
        return Tier.HEAVY

    try:
        if "```json" in raw:
            raw = raw.split("```json", 1)[1].split("```", 1)[0].strip()
        elif "```" in raw:
            raw = raw.split("```", 1)[1].split("```", 1)[0].strip()
        parsed = json.loads(raw)
        tier_value = parsed.get("tier", "heavy")
        return Tier(tier_value)
    except (json.JSONDecodeError, ValueError):
        logger.warning("Router returned unparseable response: %s", raw[:200])
        return Tier.HEAVY
