"""ExtractionPipeline: 5-stage auto memory extraction from conversations.

Stage 1: LLM Fact Extraction (with confidence evaluation)
         Phase 8.5: Also extracts entities and relations for Knowledge Graph.
Stage 2: Embedding & Similarity Search
Stage 2.5: MinHash Pre-filter (near-duplicate → auto-NOOP without LLM)
Stage 3: LLM Conflict Resolution (batched)
Stage 4: ZTM Trust Assignment (handled by caller)
Stage 5: Batch Commit (handled by caller via add_memory)

This module handles Stages 1-3 (including 2.5). Stages 4-5 are handled by
the handler (PersistentHandler) that calls this pipeline.
"""
from __future__ import annotations

import hashlib
import json
import logging
import re
from typing import Any, Optional

from .entities import VALID_ENTITY_TYPES, normalize_predicate
from .prompts import (
    BATCH_CONFLICT_PROMPT,
    CONFLICT_SIMILARITY_THRESHOLD,
    MINHASH_NOOP_THRESHOLD,
    MINHASH_NUM_PERM,
    UNIFIED_EXTRACTION_PROMPT,
    UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT,
    format_conversation,
)

logger = logging.getLogger(__name__)


# ==============================================================
# MinHash utilities (pure Python, no external dependencies)
# ==============================================================


def _char_ngrams(text: str, n: int = 3) -> set[str]:
    """Generate character n-gram shingles from text.

    Args:
        text: Input text (will be lowercased and stripped).
        n: N-gram size (default: 3).

    Returns:
        Set of character n-gram strings.
    """
    text = text.strip().lower()
    if len(text) < n:
        return {text} if text else set()
    return {text[i:i + n] for i in range(len(text) - n + 1)}


def _minhash_signature(shingles: set[str], num_perm: int = MINHASH_NUM_PERM) -> list[int]:
    """Generate MinHash signature using hash function family.

    Uses MD5 with different seeds to simulate independent hash functions.

    Args:
        shingles: Set of shingle strings.
        num_perm: Number of hash permutations (default: 128).

    Returns:
        List of minimum hash values (one per permutation).
    """
    if not shingles:
        return [0] * num_perm
    sig = [float("inf")] * num_perm
    for shingle in shingles:
        for i in range(num_perm):
            h = int(hashlib.md5(f"{i}:{shingle}".encode()).hexdigest(), 16)
            if h < sig[i]:
                sig[i] = h
    return sig


def _minhash_jaccard(sig_a: list[int], sig_b: list[int]) -> float:
    """Estimate Jaccard similarity from two MinHash signatures.

    Args:
        sig_a: MinHash signature of set A.
        sig_b: MinHash signature of set B.

    Returns:
        Estimated Jaccard similarity (0.0 to 1.0).
    """
    if not sig_a or not sig_b or len(sig_a) != len(sig_b):
        return 0.0
    return sum(a == b for a, b in zip(sig_a, sig_b)) / len(sig_a)


class ExtractionPipeline:
    """Auto memory extraction pipeline.

    Uses AIClient for LLM operations and VectorStore for similarity search.
    Gracefully degrades when AI client returns stub responses or vectorstore
    is unavailable.
    """

    def __init__(self, ai_client: Any, vectorstore: Any = None) -> None:
        """Initialize pipeline.

        Args:
            ai_client: AIClient instance (embed + llm_judge). Required.
            vectorstore: VectorStore instance for similarity search. Optional.
                         When None, Stage 2-3 are skipped (all facts → ADD).
        """
        self._ai = ai_client
        self._vs = vectorstore

    def run(
        self,
        messages: list[dict],
        user_id: str | None = None,
        custom_extraction_prompt: str | None = None,
        auto_resolve_conflicts: bool = True,
        extract_entities: bool = True,
    ) -> dict:
        """Execute the extraction pipeline (Stages 1-3).

        Args:
            messages: Conversation messages [{"role": "...", "content": "..."}]
            user_id: User ID for scoping similarity search
            custom_extraction_prompt: Override default extraction prompt.
                                     Must contain {conversation} placeholder.
            auto_resolve_conflicts: If True, use LLM for conflict resolution.
                                   If False, all similar matches → ADD.
            extract_entities: If True, also extract entities and relations
                             for Knowledge Graph (Phase 8.5). Default True.

        Returns:
            {
                "facts": [
                    {
                        "content": str,
                        "confidence": "high|medium|low",
                        "action": "ADD|NOOP|CONTRADICT",
                        "reason": str,
                        "conflict_with": str | None,  # existing memory_id
                    }
                ],
                "entities": [{"name": str, "type": str}],       # Phase 8.5
                "relations": [{"subject": str, "predicate": str, "object": str}],  # Phase 8.5
                "llm_calls": int,
            }
        """
        llm_calls = 0

        # ----------------------------------------------------------
        # Stage 1: LLM Fact Extraction (+ entities/relations)
        # ----------------------------------------------------------
        facts, entities, relations, stage1_calls, extraction_failed = self._extract_facts(
            messages, custom_extraction_prompt, extract_entities=extract_entities,
        )
        llm_calls += stage1_calls

        if extraction_failed:
            return {
                "facts": [], "entities": [], "relations": [],
                "llm_calls": llm_calls,
                "extraction_failed": True,
            }

        if not facts and not entities:
            return {
                "facts": [], "entities": [], "relations": [],
                "llm_calls": llm_calls,
            }

        # ----------------------------------------------------------
        # Stage 2: Embedding & Similarity Search (batch-optimized)
        # ----------------------------------------------------------
        conflict_pairs = []
        if self._vs is not None and auto_resolve_conflicts and facts:
            # Batch embed all facts at once if supported
            fact_texts = [f["content"] for f in facts]
            fact_vecs = self._batch_embed(fact_texts)

            for i, fact in enumerate(facts):
                vec = fact_vecs[i] if i < len(fact_vecs) else None
                similar = self._find_similar_with_vec(vec, user_id) if vec else []
                if similar:
                    # Take the most similar match
                    best = similar[0]
                    payload = best.get("payload", {})
                    conflict_pairs.append({
                        "index": i,
                        "new_fact": fact["content"],
                        "existing_content": payload.get("content", ""),
                        "existing_id": best.get("id", ""),
                        "existing_created_at": payload.get("created_at", ""),
                        "similarity": best.get("score", 0.0),
                    })

        # ----------------------------------------------------------
        # Stage 2.5: MinHash Pre-filter (near-duplicate → auto-NOOP)
        # ----------------------------------------------------------
        minhash_noops: set[int] = set()
        remaining_pairs = []
        for pair in conflict_pairs:
            existing_content = pair.get("existing_content", "")
            if not existing_content:
                # No content in payload (legacy data) → pass through to LLM
                remaining_pairs.append(pair)
                continue
            new_shingles = _char_ngrams(pair["new_fact"])
            existing_shingles = _char_ngrams(existing_content)
            if not new_shingles or not existing_shingles:
                remaining_pairs.append(pair)
                continue
            new_sig = _minhash_signature(new_shingles)
            existing_sig = _minhash_signature(existing_shingles)
            jaccard = _minhash_jaccard(new_sig, existing_sig)
            if jaccard >= MINHASH_NOOP_THRESHOLD:
                minhash_noops.add(pair["index"])
                logger.debug(
                    "MinHash auto-NOOP: '%s' ≈ '%s' (jaccard=%.3f)",
                    pair["new_fact"][:50], existing_content[:50], jaccard,
                )
            else:
                remaining_pairs.append(pair)

        # ----------------------------------------------------------
        # Stage 3: LLM Conflict Resolution
        # ----------------------------------------------------------
        resolutions: dict[int, dict] = {}

        # Add MinHash auto-NOOP resolutions first (preserve existing_id)
        pair_by_index = {p["index"]: p for p in conflict_pairs}
        for idx in minhash_noops:
            resolutions[idx] = {
                "index": idx,
                "action": "NOOP",
                "reason": "near-duplicate (MinHash auto-detected)",
                "existing_id": pair_by_index[idx].get("existing_id"),
            }

        # Resolve remaining pairs with LLM
        if remaining_pairs and auto_resolve_conflicts:
            resolved, stage3_calls = self._resolve_conflicts(remaining_pairs)
            llm_calls += stage3_calls
            for r in resolved:
                idx = r.get("index", -1)
                if 0 <= idx < len(facts):
                    resolutions[idx] = r

        # ----------------------------------------------------------
        # Build result
        # ----------------------------------------------------------
        result_facts = []
        for i, fact in enumerate(facts):
            if i in resolutions:
                res = resolutions[i]
                pair = pair_by_index.get(i, {})
                existing_id = res.get("existing_id") or pair.get("existing_id")
                entry = {
                    "content": fact["content"],
                    "confidence": fact.get("confidence", "medium"),
                    "action": res.get("action", "ADD"),
                    "reason": res.get("reason", ""),
                    "conflict_with": existing_id,
                }
                if res.get("action") == "CONTRADICT":
                    entry["winner"] = res.get("winner", "unknown")
                    entry["loser_confidence"] = res.get("loser_confidence", "medium")
                result_facts.append(entry)
            else:
                result_facts.append({
                    "content": fact["content"],
                    "confidence": fact.get("confidence", "medium"),
                    "action": "ADD",
                    "reason": "new fact (no similar memory found)",
                    "conflict_with": None,
                })

        return {
            "facts": result_facts,
            "entities": entities,
            "relations": relations,
            "llm_calls": llm_calls,
        }

    # ==============================================================
    # Stage 1: Fact Extraction
    # ==============================================================

    def _extract_facts(
        self,
        messages: list[dict],
        custom_prompt: str | None = None,
        extract_entities: bool = True,
    ) -> tuple[list[dict], list[dict], list[dict], int, bool]:
        """Extract atomic facts (and optionally entities/relations) using LLM.

        Returns:
            (facts_list, entities_list, relations_list, llm_call_count, extraction_failed)
        """
        conversation_text = format_conversation(messages)

        if custom_prompt and "{conversation}" in custom_prompt:
            prompt = custom_prompt.replace("{conversation}", conversation_text)
        elif extract_entities:
            prompt = UNIFIED_EXTRACTION_WITH_ENTITIES_PROMPT.replace(
                "{conversation}", conversation_text,
            )
        else:
            prompt = UNIFIED_EXTRACTION_PROMPT.replace("{conversation}", conversation_text)

        try:
            response = self._ai.llm_judge(prompt)
        except Exception as e:
            logger.warning("LLM extraction failed: %s", e)
            return [], [], [], 1, True

        parsed = _parse_json_safe(response)
        if parsed is None:
            logger.warning("Failed to parse LLM extraction response")
            return [], [], [], 1, True

        # --- Parse facts (Phase 8) ---
        raw_facts = parsed.get("facts", [])
        facts = []
        for item in raw_facts:
            if isinstance(item, str):
                facts.append({"content": item, "confidence": "medium"})
            elif isinstance(item, dict) and "content" in item:
                confidence = item.get("confidence", "medium")
                if confidence not in ("high", "medium", "low"):
                    confidence = "medium"
                facts.append({"content": item["content"], "confidence": confidence})

        # --- Parse entities (Phase 8.5) ---
        entities: list[dict] = []
        if extract_entities:
            raw_entities = parsed.get("entities", [])
            for ent in raw_entities:
                if isinstance(ent, dict) and "name" in ent:
                    ent_type = ent.get("type", "other")
                    if ent_type not in VALID_ENTITY_TYPES:
                        ent_type = "other"
                    entities.append({
                        "name": ent["name"].strip(),
                        "type": ent_type,
                    })

        # --- Parse relations (Phase 8.5) ---
        relations: list[dict] = []
        if extract_entities:
            raw_relations = parsed.get("relations", [])
            entity_names = {e["name"].lower() for e in entities}
            for rel in raw_relations:
                if not isinstance(rel, dict):
                    continue
                subj = rel.get("subject", "")
                pred = rel.get("predicate", "")
                obj = rel.get("object", "")
                if not all((subj, pred, obj)):
                    continue
                # Validate that subject and object are in extracted entities
                # (relaxed: allow even if not in list for robustness)
                relations.append({
                    "subject": subj.strip(),
                    "predicate": normalize_predicate(pred),
                    "object": obj.strip(),
                })

        return facts, entities, relations, 1, False

    # ==============================================================
    # Stage 2: Similarity Search (batch-optimized)
    # ==============================================================

    def _batch_embed(self, texts: list[str]) -> list[list[float] | None]:
        """Batch embed texts, using batch_embed if available, else individual calls."""
        if not texts:
            return []
        # Use batch_embed if the AI client supports it
        if hasattr(self._ai, "batch_embed"):
            try:
                return self._ai.batch_embed(texts)
            except Exception as e:
                logger.warning("Batch embedding failed, falling back to individual: %s", e)
        # Fallback: individual embed calls
        results: list[list[float] | None] = []
        for text in texts:
            try:
                results.append(self._ai.embed(text))
            except Exception as e:
                logger.warning("Embedding failed for similarity search: %s", e)
                results.append(None)
        return results

    def _find_similar_with_vec(
        self,
        vec: list[float],
        user_id: str | None = None,
    ) -> list[dict]:
        """Find existing memories similar using a precomputed embedding vector."""
        if self._vs is None or vec is None:
            return []

        try:
            vs_filters: dict[str, Any] = {
                "memory_type": "semantic",
            }
            if user_id is not None:
                vs_filters["user_id"] = user_id
            results = self._vs.search(
                query_vector=vec,
                limit=5,
                filters=vs_filters,
            )
        except Exception as e:
            logger.warning("VectorStore search failed: %s", e)
            return []

        return [r for r in results if r.get("score", 0.0) >= CONFLICT_SIMILARITY_THRESHOLD]

    def _find_similar(
        self,
        fact_content: str,
        user_id: str | None = None,
    ) -> list[dict]:
        """Find existing memories similar to a fact (single-embed fallback)."""
        if self._vs is None:
            return []

        try:
            vec = self._ai.embed(fact_content)
        except Exception as e:
            logger.warning("Embedding failed for similarity search: %s", e)
            return []

        return self._find_similar_with_vec(vec, user_id)

    # ==============================================================
    # Stage 3: Conflict Resolution
    # ==============================================================

    def _resolve_conflicts(
        self,
        pairs: list[dict],
    ) -> tuple[list[dict], int]:
        """Resolve conflicts between new facts and existing memories using LLM.

        Args:
            pairs: List of conflict pairs with index, new_fact, existing_content, etc.

        Returns:
            (resolutions, llm_call_count)
            resolutions: [{"index": 0, "action": "ADD|CONTRADICT|NOOP", "reason": "..."}]
        """
        if not pairs:
            return [], 0

        # Build pairs JSON for the prompt
        pairs_for_prompt = []
        for p in pairs:
            pairs_for_prompt.append({
                "index": p["index"],
                "new_fact": p["new_fact"],
                "existing_memory": p["existing_content"],
                "existing_created_at": p.get("existing_created_at", "unknown"),
            })

        pairs_json = json.dumps(pairs_for_prompt, ensure_ascii=False, indent=2)
        prompt = BATCH_CONFLICT_PROMPT.replace("{pairs_json}", pairs_json)

        try:
            response = self._ai.llm_judge(prompt)
        except Exception as e:
            logger.warning("LLM conflict resolution failed: %s", e)
            # Fall back: conservative NOOP to avoid duplicates
            return [{"index": p["index"], "action": "NOOP", "reason": "conflict resolution failed"} for p in pairs], 1

        parsed = _parse_json_safe(response)
        if parsed is None:
            return [{"index": p["index"], "action": "NOOP", "reason": "conflict resolution parse error"} for p in pairs], 1

        results = parsed.get("results", [])

        # Validate actions
        valid_actions = {"ADD", "CONTRADICT", "NOOP"}
        valid_winners = {"new", "existing", "unknown"}
        valid_confidences = {"high", "medium", "low"}
        validated = []
        for r in results:
            action = r.get("action", "NOOP").upper()
            if action not in valid_actions:
                action = "NOOP"
            entry = {
                "index": r.get("index", 0),
                "action": action,
                "reason": r.get("reason", ""),
            }
            if action == "CONTRADICT":
                winner = r.get("winner", "unknown")
                if winner not in valid_winners:
                    winner = "unknown"
                loser_confidence = r.get("loser_confidence", "medium")
                if loser_confidence not in valid_confidences:
                    loser_confidence = "medium"
                entry["winner"] = winner
                entry["loser_confidence"] = loser_confidence
            validated.append(entry)

        return validated, 1


# ==============================================================
# JSON parsing utility
# ==============================================================

def _parse_json_safe(text: str) -> Optional[dict]:
    """Parse JSON from LLM response with error recovery.

    Handles:
    - Clean JSON
    - JSON wrapped in markdown code fences (```json ... ```)
    - Leading/trailing whitespace
    - Returns None on failure
    """
    if not text or not text.strip():
        return None

    text = text.strip()

    # Try direct parse first
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Remove markdown code fences
    # Pattern: ```json\n{...}\n``` or ```\n{...}\n```
    fence_pattern = re.compile(r"```(?:json)?\s*\n?(.*?)\n?\s*```", re.DOTALL)
    match = fence_pattern.search(text)
    if match:
        try:
            return json.loads(match.group(1).strip())
        except json.JSONDecodeError:
            pass

    # Try to find JSON object in the text
    brace_start = text.find("{")
    brace_end = text.rfind("}")
    if brace_start >= 0 and brace_end > brace_start:
        try:
            return json.loads(text[brace_start:brace_end + 1])
        except json.JSONDecodeError:
            pass

    return None
