"""
Episode merging to prevent redundant episode accumulation.

Detects semantically similar episodes and merges them when appropriate.
"""

import json

from pydantic import BaseModel

from memv.models import Episode
from memv.processing.prompts import merge_content_prompt, merge_decision_prompt
from memv.protocols import EmbeddingClient, LLMClient


class MergeDecision(BaseModel):
    """Result of merge decision check."""

    should_merge: bool
    reason: str | None = None


class MergeCandidate(BaseModel):
    """A pair of episodes that might be merged."""

    episode1: Episode
    episode2: Episode
    similarity: float


class EpisodeMerger:
    """
    Merges semantically similar episodes to prevent redundancy.

    Flow:
    1. Find candidate pairs via embedding similarity
    2. LLM verifies if merge is appropriate
    3. Generate merged content (title + content)
    4. Return merged episode

    Args:
        llm_client: LLM for merge decisions and content generation
        embedding_client: For computing episode similarities
        similarity_threshold: Min cosine similarity to consider merging (default 0.85)
    """

    def __init__(
        self,
        llm_client: LLMClient,
        embedding_client: EmbeddingClient,
        similarity_threshold: float = 0.85,
    ):
        self.llm = llm_client
        self.embedder = embedding_client
        self.similarity_threshold = similarity_threshold

    async def find_merge_candidates(
        self,
        new_episode: Episode,
        existing_episodes: list[Episode],
    ) -> list[MergeCandidate]:
        """
        Find episodes that might be merged with the new episode.

        Args:
            new_episode: Newly created episode
            existing_episodes: Episodes to compare against

        Returns:
            List of merge candidates sorted by similarity (highest first)
        """
        if not existing_episodes:
            return []

        # Get embedding for new episode
        new_text = f"{new_episode.title} {new_episode.content}"
        new_embedding = await self.embedder.embed(new_text)

        # Get embeddings for existing episodes
        existing_texts = [f"{ep.title} {ep.content}" for ep in existing_episodes]
        existing_embeddings = await self.embedder.embed_batch(existing_texts)

        # Find similar episodes
        candidates: list[MergeCandidate] = []
        for ep, embedding in zip(existing_episodes, existing_embeddings, strict=True):
            similarity = self._cosine_similarity(new_embedding, embedding)
            if similarity >= self.similarity_threshold:
                candidates.append(
                    MergeCandidate(
                        episode1=new_episode,
                        episode2=ep,
                        similarity=similarity,
                    )
                )

        # Sort by similarity descending
        candidates.sort(key=lambda c: c.similarity, reverse=True)
        return candidates

    async def should_merge(self, candidate: MergeCandidate) -> MergeDecision:
        """
        Use LLM to verify if two episodes should actually be merged.

        Embedding similarity alone can have false positives. LLM provides
        semantic verification.
        """
        prompt = merge_decision_prompt(
            episode1_title=candidate.episode1.title,
            episode1_content=candidate.episode1.content,
            episode2_title=candidate.episode2.title,
            episode2_content=candidate.episode2.content,
        )
        response = await self.llm.generate(prompt)
        return self._parse_decision(response)

    async def merge(self, episode1: Episode, episode2: Episode) -> Episode:
        """
        Merge two episodes into one.

        The merged episode:
        - Combines original_messages from both
        - Has new title/content generated by LLM
        - Uses earliest start_time and latest end_time
        - Preserves user_id (must be same for both)

        Args:
            episode1: First episode
            episode2: Second episode

        Returns:
            New merged episode
        """
        if episode1.user_id != episode2.user_id:
            raise ValueError("Cannot merge episodes from different users")

        # Generate merged content
        prompt = merge_content_prompt(
            episode1_title=episode1.title,
            episode1_content=episode1.content,
            episode2_title=episode2.title,
            episode2_content=episode2.content,
        )
        response = await self.llm.generate(prompt)
        title, content = self._parse_merge_content(response)

        # Combine original messages (preserve order)
        combined_messages = episode1.original_messages + episode2.original_messages

        return Episode(
            user_id=episode1.user_id,
            title=title,
            content=content,
            original_messages=combined_messages,
            start_time=min(episode1.start_time, episode2.start_time),
            end_time=max(episode1.end_time, episode2.end_time),
        )

    async def merge_if_appropriate(
        self,
        new_episode: Episode,
        existing_episodes: list[Episode],
    ) -> tuple[Episode, Episode | None]:
        """
        Check if new episode should be merged with an existing one.

        Args:
            new_episode: Newly created episode
            existing_episodes: Episodes to compare against

        Returns:
            (result_episode, merged_with) - merged_with is None if no merge happened
        """
        candidates = await self.find_merge_candidates(new_episode, existing_episodes)

        for candidate in candidates:
            decision = await self.should_merge(candidate)
            if decision.should_merge:
                merged = await self.merge(candidate.episode1, candidate.episode2)
                return merged, candidate.episode2

        return new_episode, None

    def _cosine_similarity(self, vec1: list[float], vec2: list[float]) -> float:
        """Compute cosine similarity between two vectors."""
        dot_product = sum(a * b for a, b in zip(vec1, vec2, strict=True))
        norm1 = sum(a * a for a in vec1) ** 0.5
        norm2 = sum(b * b for b in vec2) ** 0.5
        if norm1 == 0 or norm2 == 0:
            return 0.0
        return dot_product / (norm1 * norm2)

    def _parse_decision(self, response: str) -> MergeDecision:
        """Parse LLM response into MergeDecision."""
        try:
            # Handle markdown code blocks
            if "```json" in response:
                start = response.index("```json") + 7
                end = response.index("```", start)
                response = response[start:end]
            elif "```" in response:
                start = response.index("```") + 3
                end = response.index("```", start)
                response = response[start:end]

            data = json.loads(response.strip())
            return MergeDecision(
                should_merge=bool(data.get("should_merge", False)),
                reason=data.get("reason"),
            )
        except (json.JSONDecodeError, ValueError, KeyError):
            # Default to no merge on parse failure
            return MergeDecision(should_merge=False, reason="Failed to parse LLM response")

    def _parse_merge_content(self, response: str) -> tuple[str, str]:
        """Parse merged episode content from LLM response."""
        try:
            # Handle markdown code blocks
            if "```json" in response:
                start = response.index("```json") + 7
                end = response.index("```", start)
                response = response[start:end]
            elif "```" in response:
                start = response.index("```") + 3
                end = response.index("```", start)
                response = response[start:end]

            data = json.loads(response.strip())
            title = data.get("title", "Merged Episode")
            content = data.get("content", "")

            if not content:
                raise ValueError("Empty content")

            return title, content

        except (json.JSONDecodeError, ValueError, KeyError):
            # Fallback
            return "Merged Episode", response.strip()
