"""Extraction service -- orchestrates LLM calls for structured extraction."""

from __future__ import annotations

import asyncio
import json
import logging
import time

from ai_core.caching import parse_json_response
from ai_core.models import AIMessage, Role
from ai_core.router import AIClient
from errors import ValidationError

from extraction.models import ExtractionConfig, ExtractionResult, ExtractionSchema
from extraction.prompts import (
    build_classification_prompt,
    build_extraction_prompt,
    build_user_prompt,
)

try:
    import structlog

    logger = structlog.get_logger("extraction")
except ImportError:
    logger = logging.getLogger("extraction")


class ExtractionService:
    """Domain-agnostic structured data extraction powered by LLMs.

    Instantiate with an :class:`AIClient` and optional :class:`ExtractionConfig`,
    then call ``extract()``, ``extract_batch()``, or ``classify()``.
    """

    def __init__(self, ai_client: AIClient, config: ExtractionConfig | None = None) -> None:
        self._ai = ai_client
        self._config = config or ExtractionConfig()

    async def extract(
        self,
        text: str,
        schema: ExtractionSchema,
        *,
        model: str | None = None,
    ) -> ExtractionResult:
        """Extract structured data from text according to schema.

        Args:
            text: The raw text to extract from.
            schema: The extraction schema describing which fields to extract.
            model: LLM model override.  Falls back to ``config.default_model``.

        Returns:
            An :class:`ExtractionResult` with the extracted key-value pairs.

        Raises:
            ValidationError: If the LLM response cannot be parsed as valid JSON.
        """
        target_model = model or self._config.default_model
        system_prompt = build_extraction_prompt(schema)
        user_prompt = build_user_prompt(text)

        messages = [
            AIMessage(role=Role.SYSTEM, content=system_prompt),
            AIMessage(role=Role.USER, content=user_prompt),
        ]

        start = time.monotonic()
        response = await self._ai.chat(
            messages,
            model=target_model,
            temperature=self._config.temperature,
            max_tokens=self._config.max_tokens,
        )
        elapsed_ms = int((time.monotonic() - start) * 1000)

        try:
            data = parse_json_response(response.text)
        except json.JSONDecodeError as exc:
            logger.warning(
                "Failed to parse extraction response as JSON: model=%s schema=%s error=%s",
                target_model,
                schema.name,
                exc,
            )
            raise ValidationError(
                f"LLM returned invalid JSON for schema '{schema.name}': {exc}"
            ) from exc

        if not isinstance(data, dict):
            raise ValidationError(f"Expected JSON object from LLM, got {type(data).__name__}")

        warnings = self._validate_against_schema(data, schema)

        return ExtractionResult(
            schema_name=schema.name,
            data=data,
            model_used=response.model,
            tokens_used=response.tokens_used.input_tokens + response.tokens_used.output_tokens,
            cost_usd=response.tokens_used.cost_usd,
            processing_time_ms=elapsed_ms,
            warnings=warnings,
        )

    async def extract_batch(
        self,
        texts: list[str],
        schema: ExtractionSchema,
        *,
        model: str | None = None,
        concurrency: int = 5,
    ) -> list[ExtractionResult]:
        """Extract from multiple texts concurrently.

        Args:
            texts: List of texts to extract from.
            schema: The extraction schema to apply to each text.
            model: LLM model override.
            concurrency: Maximum number of concurrent extractions.

        Returns:
            A list of :class:`ExtractionResult` in the same order as ``texts``.
        """
        semaphore = asyncio.Semaphore(concurrency)

        async def _extract_one(text: str) -> ExtractionResult:
            async with semaphore:
                return await self.extract(text, schema, model=model)

        return list(await asyncio.gather(*[_extract_one(t) for t in texts]))

    async def classify(
        self,
        text: str,
        categories: list[str],
        *,
        model: str | None = None,
    ) -> dict:
        """Classify text into one of the given categories.

        Args:
            text: The text to classify.
            categories: List of valid category labels.
            model: LLM model override.

        Returns:
            A dict with ``category`` and ``confidence`` keys.

        Raises:
            ValidationError: If the LLM response cannot be parsed or the
                returned category is not in the provided list.
        """
        target_model = model or self._config.default_model
        system_prompt = build_classification_prompt(categories)
        user_prompt = f"Classify the following text:\n\n{text}"

        messages = [
            AIMessage(role=Role.SYSTEM, content=system_prompt),
            AIMessage(role=Role.USER, content=user_prompt),
        ]

        response = await self._ai.chat(
            messages,
            model=target_model,
            temperature=0.0,
            max_tokens=256,
        )

        try:
            data = parse_json_response(response.text)
        except json.JSONDecodeError as exc:
            raise ValidationError(f"LLM returned invalid JSON for classification: {exc}") from exc

        if not isinstance(data, dict):
            raise ValidationError(f"Expected JSON object from LLM, got {type(data).__name__}")

        category = data.get("category", "")
        if category not in categories:
            logger.warning(
                "LLM returned unknown category: %s (valid: %s)",
                category,
                categories,
            )

        return data

    @staticmethod
    def _validate_against_schema(
        data: dict,
        schema: ExtractionSchema,
    ) -> list[str]:
        """Check extracted data against the schema, returning warnings.

        This is a soft validation -- it does not raise but collects warnings
        for missing required fields or unexpected keys.
        """
        warnings: list[str] = []
        field_names = {f.name for f in schema.fields}

        for field in schema.fields:
            if field.required and field.name not in data:
                warnings.append(f"Required field '{field.name}' missing from extraction")

        extra_keys = set(data.keys()) - field_names
        if extra_keys:
            warnings.append(f"Unexpected keys in extraction: {sorted(extra_keys)}")

        return warnings
