"""Gist extraction from episode clusters."""

import json
from dataclasses import dataclass
from typing import Any

from ..utils.llm import LLMClient
from .clusterer import EpisodeCluster

GIST_EXTRACTION_PROMPT = """Analyze these related memories and extract the key semantic information.

MEMORIES (from a conversation with a user):
{memories}

COMMON THEMES: {themes}

Extract:
1. The main fact or pattern these memories represent
2. The confidence level (how consistent/certain the info is)
3. The type: one of:
   - "fact" (definite info)
   - "preference" (user likes/dislikes)
   - "pattern" (behavioral tendency)
   - "summary" (general synopsis)
   - "goal" (something the user is working toward or trying to achieve)
   - "value" (something the user considers important or prioritizes)
   - "state" (a current condition or situation the user is in)
   - "causal" (a reason or explanation for user behavior)
   - "policy" (a personal rule the user follows, e.g. "I never...", "I always...")
4. A structured representation if possible (subject, predicate, value)

Return JSON:
{{
  "gist": "User prefers vegetarian food",
  "type": "preference",
  "confidence": 0.9,
  "subject": "user",
  "predicate": "food_preference",
  "value": "vegetarian",
  "key": "user:preference:food"
}}

Rules:
- Combine information across memories to get the core meaning
- Don't include episodic details (times, specific conversations)
- Focus on durable, generalizable information
- Higher confidence if multiple memories support the same conclusion
- Use "goal"/"value"/"state"/"causal"/"policy" types when memories express constraints, commitments, or conditions that should govern future behavior"""

# Constraint-signal words to preserve when gist_type is cognitive
_COGNITIVE_GIST_TYPES = {"goal", "value", "state", "causal", "policy"}
_CONSTRAINT_SIGNAL_WORDS = frozenset(
    {
        "never",
        "always",
        "must",
        "avoid",
        "allergic",
        "shouldn't",
        "can't",
        "won't",
        "refuse",
        "should",
        "need",
        "have to",
        "trying",
        "value",
        "important",
        "because",
        "reason",
        "policy",
        "rule",
    }
)


@dataclass
class ExtractedGist:
    """Extracted semantic gist from a cluster."""

    text: str
    gist_type: str  # "fact", "preference", "pattern", "summary", "goal", "value", "state", "causal", "policy"
    confidence: float
    supporting_episode_ids: list[str]

    key: str | None = None
    subject: str | None = None
    predicate: str | None = None
    value: Any | None = None


class GistExtractor:
    """Extracts semantic gist from episode clusters. Uses heuristic when llm_client is None."""

    def __init__(self, llm_client: LLMClient | None):
        self.llm = llm_client

    def _heuristic_extract_gist(self, cluster: EpisodeCluster) -> list[ExtractedGist]:
        """Heuristic fallback: truncation + keyword-based gist_type when LLM disabled."""
        if not cluster.episodes:
            return []
        combined = " ".join(ep.text for ep in cluster.episodes).lower()
        gist_text = combined[:100].strip()
        if len(combined) > 100:
            gist_text = (
                gist_text.rsplit(maxsplit=1)[0] + "..." if " " in gist_text else gist_text + "..."
            )
        gist_type = "summary"
        if any(w in combined for w in ("prefer", "like", "love", "enjoy", "hate", "dislike")):
            gist_type = "preference"
        elif any(w in combined for w in ("never", "always", "must", "should")):
            gist_type = "policy"
        elif any(w in combined for w in ("goal", "trying", "working toward")):
            gist_type = "goal"
        elif any(w in combined for w in ("value", "important", "believe")):
            gist_type = "value"
        return [
            ExtractedGist(
                text=gist_text or "Cluster summary",
                gist_type=gist_type,
                confidence=0.5,
                supporting_episode_ids=[str(ep.id) for ep in cluster.episodes],
            )
        ]

    async def extract_gist(self, cluster: EpisodeCluster) -> list[ExtractedGist]:
        """Extract gists from a single cluster. Uses heuristic when LLM is None."""
        if not cluster.episodes:
            return []

        if self.llm is None:
            return self._heuristic_extract_gist(cluster)

        memory_texts = []
        for i, ep in enumerate(cluster.episodes[:10], 1):
            mem_type = ep.type.value if hasattr(ep.type, "value") else str(ep.type)
            memory_texts.append(f"{i}. [{mem_type}] {ep.text}")

        memories_str = "\n".join(memory_texts)
        themes_str = (
            ", ".join(cluster.common_entities) if cluster.common_entities else "none identified"
        )

        prompt = GIST_EXTRACTION_PROMPT.format(
            memories=memories_str,
            themes=themes_str,
        )

        try:
            response = await self.llm.complete(prompt, temperature=0.0)
            # Strip markdown code block if present
            raw = response.strip()
            if raw.startswith("```"):
                lines = raw.split("\n")
                if lines[0].startswith("```"):
                    lines = lines[1:]
                if lines and lines[-1].strip() == "```":
                    lines = lines[:-1]
                raw = "\n".join(lines)
            data = json.loads(raw)

            gists_data = data if isinstance(data, list) else [data]

            gists = []
            for gd in gists_data:
                gist_text = gd.get("gist", "").strip()
                gist_type = gd.get("type", "summary")
                if not gist_text:
                    continue
                # Optional validation: cognitive gists should preserve constraint-signal words
                if gist_type in _COGNITIVE_GIST_TYPES:
                    cluster_text = " ".join(ep.text for ep in cluster.episodes).lower()
                    gist_lower = gist_text.lower()
                    signals_in_cluster = [w for w in _CONSTRAINT_SIGNAL_WORDS if w in cluster_text]
                    if signals_in_cluster and not any(w in gist_lower for w in signals_in_cluster):
                        continue  # Drop gist that lost constraint semantics
                conf = float(gd.get("confidence", 0.7)) * cluster.avg_confidence
                gists.append(
                    ExtractedGist(
                        text=gist_text,
                        gist_type=gist_type,
                        confidence=conf,
                        supporting_episode_ids=[str(ep.id) for ep in cluster.episodes],
                        key=gd.get("key"),
                        subject=gd.get("subject"),
                        predicate=gd.get("predicate"),
                        value=gd.get("value"),
                    )
                )
            return gists

        except (json.JSONDecodeError, KeyError, TypeError):
            return [
                ExtractedGist(
                    text=self._simple_summary(cluster),
                    gist_type="summary",
                    confidence=cluster.avg_confidence * 0.5,
                    supporting_episode_ids=[str(ep.id) for ep in cluster.episodes],
                )
            ]

    async def extract_from_clusters(self, clusters: list[EpisodeCluster]) -> list[ExtractedGist]:
        """Extract gists from all clusters."""
        import asyncio

        all_gists = []
        results = await asyncio.gather(*[self.extract_gist(c) for c in clusters])
        for gist_list in results:
            all_gists.extend(gist_list)
        return all_gists

    def _simple_summary(self, cluster: EpisodeCluster) -> str:
        if cluster.common_entities:
            return f"User discussed: {', '.join(cluster.common_entities[:3])}"
        if cluster.episodes:
            return cluster.episodes[0].text[:100]
        return "Cluster summary"
