"""DiscoveryExtractor: turns five narrative responses into a structured
enneagram + Ikigai profile via the Anthropic API.

The extractor is responsible only for plumbing: building the message,
calling the model, parsing/validating the JSON it returns, and retrying
once if the first response is not valid JSON. The actual analysis lives
in the model behind ``SYSTEM_PROMPT``.
"""
import json
from typing import Any, Dict

from anthropic import AsyncAnthropic

DEFAULT_MODEL = "claude-sonnet-4-5"
MAX_TOKENS = 2000

# Ordered mapping of response keys -> the human-facing prompt label, used to
# label each section of the concatenated user message.
PROMPT_LABELS = {
    "alive": "The Alive Moment",
    "friction": "The Friction Moment",
    "pull": "The Natural Pull",
    "recognition": "The Recognition Moment",
    "future": "The Future Pull",
}

# Keys the model must return for a profile to be considered well-formed.
REQUIRED_KEYS = (
    "triad",
    "probable_type",
    "wing",
    "instinctual_variant",
    "instinctual_stack",
    "love_summary",
    "strength_summary",
    "mission_summary",
    "vocation_summary",
    "overlap_narrative",
    "confidence",
)
REQUIRED_CONFIDENCE_KEYS = ("triad", "type", "variant", "ikigai")

SYSTEM_PROMPT = """You are a skilled personality analyst trained in the Enneagram system and the Ikigai framework.
You will receive five narrative responses from a person answering open-ended reflection prompts.
Your job is to extract a structured self-discovery profile from their stories.

ENNEAGRAM EXTRACTION RULES:
- The nine types cluster into three triads based on emotional center:
  - Gut (instinctive): Types 8, 9, 1 — driven by anger, focused on control, body-based decisions
  - Heart (feeling): Types 2, 3, 4 — driven by shame, focused on image and connection
  - Head (thinking): Types 5, 6, 7 — driven by fear, focused on safety and understanding
- The Friction Moment response reveals the triad most clearly — gut types act against injustice, heart types feel exposed or unseen, head types analyze and strategize
- The Alive Moment and Recognition Moment reveal the type's core need
- The Natural Pull reveals instinctual variant: sp (self-preservation) = tasks/systems/stability, so (social) = groups/community/belonging, sx (sexual/one-to-one) = intensity/connection/depth
- The Future Pull reveals the type's idealized self and Ikigai vocation

IKIGAI EXTRACTION RULES:
- Love: what activities, topics, or experiences appear across responses with energy and enthusiasm
- Strength: what the person describes doing well or being recognized for
- Mission: what problem or need in the world their stories orbit around
- Vocation: where their strength and the world's need intersect with economic potential

CONFIDENCE RULES:
- high: strong consistent signal across 2+ responses
- medium: signal present but only in one response or partially contradicted
- low: weak or absent signal — do not guess, flag it

OUTPUT FORMAT:
Respond ONLY with valid JSON. No preamble, no explanation, no markdown fences.

{
  "triad": "gut | heart | head",
  "probable_type": 1-9,
  "wing": 1-9 (must be adjacent to probable_type),
  "instinctual_variant": "sp | so | sx",
  "instinctual_stack": "e.g. sp/so/sx",
  "love_summary": "2-3 sentence summary of what they love",
  "strength_summary": "2-3 sentence summary of what they are good at",
  "mission_summary": "2-3 sentence summary of what the world needs from them",
  "vocation_summary": "2-3 sentence summary of what they can be paid for",
  "overlap_narrative": "One paragraph (4-6 sentences) describing where their four Ikigai circles converge and how their enneagram type shapes that intersection. Write directly to the person in second person (you/your). Do not mention enneagram type numbers — describe the pattern in plain language.",
  "confidence": {
    "triad": "high | medium | low",
    "type": "high | medium | low",
    "variant": "high | medium | low",
    "ikigai": "high | medium | low"
  },
  "extraction_notes": "Optional: flag anything ambiguous, contradictory, or uncertain that the user should know"
}"""

RETRY_REMINDER = (
    "Your previous response could not be parsed as JSON. "
    "Respond ONLY with the single valid JSON object described in your "
    "instructions — no preamble, no explanation, and no markdown code fences."
)


class DiscoveryExtractionError(Exception):
    """Raised when extraction fails (API error or unparseable output)."""


class DiscoveryExtractor:
    """Extracts a self-discovery profile from narrative responses."""

    def __init__(self, api_key: str, model: str = DEFAULT_MODEL):
        if not api_key:
            raise DiscoveryExtractionError(
                "ANTHROPIC_API_KEY is not set; cannot run extraction."
            )
        self.model = model
        self.client = AsyncAnthropic(api_key=api_key)

    def _build_user_message(self, responses: Dict[str, str]) -> str:
        """Concatenate the five responses, each under its prompt heading."""
        sections = []
        for key, label in PROMPT_LABELS.items():
            text = (responses.get(key) or "").strip()
            sections.append(f"## {label}\n{text if text else '(no response)'}")
        return "\n\n".join(sections)

    async def _call_model(self, user_message: str) -> str:
        """Make a single Anthropic API call and return the raw text."""
        response = await self.client.messages.create(
            model=self.model,
            max_tokens=MAX_TOKENS,
            system=SYSTEM_PROMPT,
            messages=[{"role": "user", "content": user_message}],
        )
        return response.content[0].text

    async def extract(self, responses: Dict[str, str]) -> Dict[str, Any]:
        """Run extraction. Retries once if the first output is not valid JSON.

        Args:
            responses: dict with keys alive, friction, pull, recognition, future.

        Returns:
            Parsed profile dict matching the system-prompt schema.

        Raises:
            DiscoveryExtractionError: on API failure or repeated parse failure.
        """
        user_message = self._build_user_message(responses)

        try:
            raw = await self._call_model(user_message)
        except Exception as exc:  # noqa: BLE001 - surface any SDK/transport error
            raise DiscoveryExtractionError(
                f"Anthropic API call failed: {exc}"
            ) from exc

        try:
            return self._parse(raw)
        except (json.JSONDecodeError, ValueError):
            # One retry with an explicit JSON-only reminder appended.
            retry_message = f"{user_message}\n\n{RETRY_REMINDER}"
            try:
                raw_retry = await self._call_model(retry_message)
            except Exception as exc:  # noqa: BLE001
                raise DiscoveryExtractionError(
                    f"Anthropic API call failed on retry: {exc}"
                ) from exc

            try:
                return self._parse(raw_retry)
            except (json.JSONDecodeError, ValueError) as exc:
                raise DiscoveryExtractionError(
                    f"Model did not return valid JSON after retry: {exc}"
                ) from exc

    @staticmethod
    def _strip_fences(text: str) -> str:
        """Remove a leading/trailing markdown code fence if present."""
        stripped = text.strip()
        if stripped.startswith("```"):
            # drop the opening fence line (``` or ```json)
            newline = stripped.find("\n")
            if newline != -1:
                stripped = stripped[newline + 1 :]
            if stripped.rstrip().endswith("```"):
                stripped = stripped.rstrip()[: -len("```")]
        return stripped.strip()

    @classmethod
    def _parse(cls, raw: str) -> Dict[str, Any]:
        """Parse and validate the model's JSON output.

        Raises json.JSONDecodeError if the text is not JSON, or ValueError if
        required keys are missing — both of which trigger a retry upstream.
        """
        if not raw or not raw.strip():
            raise ValueError("empty response from model")

        data = json.loads(cls._strip_fences(raw))

        if not isinstance(data, dict):
            raise ValueError("top-level JSON value is not an object")

        missing = [k for k in REQUIRED_KEYS if k not in data]
        if missing:
            raise ValueError(f"missing required keys: {', '.join(missing)}")

        confidence = data.get("confidence")
        if not isinstance(confidence, dict):
            raise ValueError("confidence must be an object")
        missing_conf = [
            k for k in REQUIRED_CONFIDENCE_KEYS if k not in confidence
        ]
        if missing_conf:
            raise ValueError(
                f"missing confidence keys: {', '.join(missing_conf)}"
            )

        return data