"""Enrichment agent orchestrating LLM calls for entity detection.

Calls the LLM with extracted text, parses the structured response,
filters entities by confidence threshold, and handles parse errors with retry.
"""

import json
import logging

from ingestion.config import IngestionConfig
from ingestion.integrations.llm import LLMClient, LLMError

from .models import EnrichmentResult, Entity
from .prompts import STRICT_RETRY_PROMPT, build_system_prompt, build_user_prompt

logger = logging.getLogger(__name__)


class EnrichmentError(Exception):
    """Raised when enrichment fails after retries."""

    pass


class EnrichmentAgent:
    """Orchestrates LLM-based document enrichment."""

    def __init__(
        self,
        config: IngestionConfig,
        llm_client: LLMClient | None = None,
        known_projects: list[str] | None = None,
        known_people: list[str] | None = None,
    ):
        self.config = config
        self.llm = llm_client or LLMClient(config)
        self.confidence_threshold = config.confidence_threshold
        self.system_prompt = build_system_prompt(known_projects, known_people)

    def enrich(self, text: str) -> EnrichmentResult:
        """Enrich extracted text with LLM-detected entities and metadata.

        Args:
            text: The extracted document text to analyze.

        Returns:
            EnrichmentResult with entities filtered by confidence threshold.

        Raises:
            EnrichmentError: If LLM call and retry both fail.
        """
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": build_user_prompt(text)},
        ]

        # First attempt
        try:
            response_text = self.llm.complete(messages)
            result = self._parse_response(response_text)
        except (json.JSONDecodeError, ValueError) as parse_err:
            logger.warning(
                "First LLM response parse failed (%s), retrying with strict prompt",
                parse_err,
            )
            # Retry with stricter prompt
            result = self._retry_with_strict_prompt(messages, response_text)
        except LLMError as e:
            raise EnrichmentError(f"LLM call failed: {e}") from e

        # Filter entities by confidence threshold
        result.entities = self._filter_by_confidence(result.entities)

        return result

    def _retry_with_strict_prompt(
        self, original_messages: list[dict[str, str]], failed_response: str
    ) -> EnrichmentResult:
        """Retry LLM call with a stricter prompt after parse failure."""
        retry_messages = original_messages + [
            {"role": "assistant", "content": failed_response},
            {"role": "user", "content": STRICT_RETRY_PROMPT},
        ]

        try:
            response_text = self.llm.complete(retry_messages)
            return self._parse_response(response_text)
        except (json.JSONDecodeError, ValueError, LLMError) as e:
            logger.error("Retry also failed (%s), using defaults", e)
            return self._default_result()

    def _parse_response(self, response_text: str) -> EnrichmentResult:
        """Parse LLM response text into an EnrichmentResult.

        Handles responses that may be wrapped in markdown code fences.
        """
        cleaned = response_text.strip()

        # Strip markdown code fences if present
        if cleaned.startswith("```"):
            lines = cleaned.split("\n")
            # Remove first line (```json) and last line (```)
            lines = [l for l in lines[1:] if not l.strip() == "```"]
            cleaned = "\n".join(lines)

        data = json.loads(cleaned)
        return EnrichmentResult.model_validate(data)

    def _filter_by_confidence(self, entities: list[Entity]) -> list[Entity]:
        """Keep only entities meeting the confidence threshold."""
        return [e for e in entities if e.confidence >= self.confidence_threshold]

    def _default_result(self) -> EnrichmentResult:
        """Return a minimal default result when all parsing fails."""
        return EnrichmentResult(
            title="Untitled Note",
            category="inbox",
            tags=[],
            entities=[],
            action_items=[],
            summary=None,
        )