"""LLM Task Classifier — async classification for low-confidence prompts.

When keyword-based task detection has low confidence (<=0.85), this module
sends the prompt to a cheap LLM for more accurate classification. Returns
a task_type + complexity score, or None on any failure.

Dependency chain: completions -> routing -> classifier -> proxy/providers
(no circular imports — providers don't import from completions or routing)
"""

from __future__ import annotations

import asyncio
import json

import structlog
from pydantic import BaseModel, ConfigDict

from robot_resources.proxy.models import ChatCompletionRequest, Message
from robot_resources.proxy.providers.anthropic import AnthropicProvider
from robot_resources.proxy.providers.base import BaseProvider, ProviderConfig
from robot_resources.proxy.providers.google import GoogleProvider
from robot_resources.proxy.providers.openai import OpenAIProvider
from robot_resources.routing.classifier_key_client import get_classifier_key_client
from robot_resources.routing.task_detection import TaskType

logger = structlog.get_logger()

# Classifier timeout in seconds
CLASSIFIER_TIMEOUT: float = 2.0

# Max prompt length sent to classifier (chars)
CLASSIFIER_MAX_PROMPT_LENGTH: int = 2000

_PROVIDER_CLASSES: dict[str, type[BaseProvider]] = {
    "openai": OpenAIProvider,
    "anthropic": AnthropicProvider,
    "google": GoogleProvider,
}

# Valid task types for runtime validation
_VALID_TASK_TYPES: set[str] = {
    "coding",
    "analysis",
    "reasoning",
    "simple_qa",
    "creative",
    "general",
}

CLASSIFICATION_PROMPT = (
    "Classify this user prompt into a task type and complexity.\n"
    "\n"
    "Task types: coding, analysis, reasoning, simple_qa, creative, general\n"
    "Complexity: 1 (trivial) to 5 (expert-level)\n"
    "\n"
    "Examples:\n"
    '- "Write a Python sort function" -> {"task_type": "coding", "complexity": 2}\n'
    '- "Explain quantum entanglement" -> {"task_type": "reasoning", "complexity": 3}\n'
    '- "What is the capital of France?" -> {"task_type": "simple_qa", "complexity": 1}\n'
    '- "Write a sonnet about the ocean" -> {"task_type": "creative", "complexity": 3}\n'
    '- "Analyze microservices vs monolith" -> {"task_type": "analysis", "complexity": 4}\n'
    '- "Build a distributed consensus algorithm" -> {"task_type": "coding", "complexity": 5}\n'
    '- "Summarize this meeting" -> {"task_type": "general", "complexity": 2}\n'
    "\n"
    "Respond with ONLY valid JSON, no other text:\n"
    '{"task_type": "<type>", "complexity": <1-5>}\n'
    "\n"
    "User prompt: "
)


class ClassificationResult(BaseModel):
    """Result of LLM classification."""

    model_config = ConfigDict(frozen=True)

    task_type: TaskType
    complexity: int  # 1-5
    classifier_model: str


async def _get_classifier_provider_async() -> tuple[str, BaseProvider] | None:
    """Get classifier model+provider from the RR-owned Platform key.

    Resolution: Platform key → None (keyword-only).
    Never uses the user's own API keys for classification — that's our cost.

    Never raises.
    """
    try:
        client = get_classifier_key_client()
        key_info = await client.get_key()
        if key_info is not None:
            provider_class = _PROVIDER_CLASSES.get(key_info.provider)
            if provider_class is not None:
                config = ProviderConfig(
                    api_key=key_info.api_key, timeout=CLASSIFIER_TIMEOUT
                )
                provider = provider_class(config)
                model = provider.map_model(key_info.model)
                return model, provider
    except Exception:
        pass

    # No Platform key available — keyword-only classification
    return None


async def classify_with_llm(prompt: str) -> ClassificationResult | None:
    """Classify a prompt using an LLM.

    Returns ``ClassificationResult`` on success, ``None`` on any failure.
    **Never raises** — all exceptions are caught and logged.

    Args:
        prompt: The user prompt to classify.

    Returns:
        Classification result or None.
    """
    if not prompt or not prompt.strip():
        return None

    try:
        return await asyncio.wait_for(
            _classify_impl(prompt),
            timeout=CLASSIFIER_TIMEOUT,
        )
    except asyncio.TimeoutError:
        logger.warning("classifier_timeout", timeout=CLASSIFIER_TIMEOUT)
        return None
    except Exception:
        logger.warning("classifier_unexpected_error", exc_info=True)
        return None


async def _classify_impl(prompt: str) -> ClassificationResult | None:
    """Internal classification implementation."""
    provider_result = await _get_classifier_provider_async()
    if provider_result is None:
        return None

    model_name, provider = provider_result

    # Truncate long prompts
    truncated = prompt[:CLASSIFIER_MAX_PROMPT_LENGTH]

    # Build request
    request = ChatCompletionRequest(
        model=model_name,
        messages=[
            Message(role="user", content=CLASSIFICATION_PROMPT + truncated),
        ],
        temperature=0.0,
        max_tokens=50,
    )

    response = await provider.complete(request)
    text = response.choices[0].message.content.strip()
    return _parse_classification(text, model_name)


def _parse_classification(text: str, model_name: str) -> ClassificationResult | None:
    """Parse LLM response text into a ClassificationResult.

    Returns ``None`` (never raises) if the text is not valid JSON
    or doesn't contain the expected fields.
    """
    # Strip markdown code fences (e.g. ```json\n...\n```) that some models wrap around JSON
    cleaned = text.strip()
    if cleaned.startswith("```"):
        cleaned = cleaned.split("\n", 1)[-1]  # drop opening fence line
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3].strip()

    try:
        data = json.loads(cleaned)
    except json.JSONDecodeError:
        logger.warning("classifier_invalid_json", response=text[:200])
        return None

    task_type = data.get("task_type")
    if task_type not in _VALID_TASK_TYPES:
        logger.warning("classifier_invalid_task_type", task_type=task_type)
        return None

    complexity = data.get("complexity")
    if not isinstance(complexity, (int, float)):
        logger.warning("classifier_invalid_complexity", complexity=complexity)
        return None

    # Clamp complexity to 1-5
    complexity = max(1, min(5, int(complexity)))

    return ClassificationResult(
        task_type=task_type,
        complexity=complexity,
        classifier_model=model_name,
    )
