"""Fact reconciliation — Mem0-style ADD/UPDATE/NOOP/DELETE decisions.

Ported from advisor_api reconcile.py.
Compares new facts against existing ones using embedding similarity,
then asks LLM to decide what to do. All LLM calls via injected provider.
"""

import json
import logging
import uuid as uuid_mod
from collections import defaultdict
from dataclasses import dataclass
from typing import Literal

from sqlalchemy import select, text
from sqlalchemy.ext.asyncio import AsyncSession

from arandu.constants import ExtractedFact
from arandu.models import MemoryFact
from arandu.protocols import EmbeddingProvider, LLMProvider
from arandu.write.entity_resolution import ResolvedEntity
from arandu.write.extract import _strip_markdown_fences

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
SIMILARITY_THRESHOLD = 0.5  # below this, skip LLM → auto ADD


# ---------------------------------------------------------------------------
# Result type
# ---------------------------------------------------------------------------
@dataclass
class ReconciliationDecision:
    """LLM decision on how a new fact relates to existing knowledge.

    Attributes:
        fact_text: The new fact being reconciled.
        subject: Entity key the fact belongs to.
        action: Reconciliation action — ADD (new), UPDATE (supersedes), NOOP (duplicate), DELETE (retract).
        matched_fact_id: ID of existing fact being updated/superseded, if any.
        matched_fact_text: Text of the matched existing fact, if any.
        confidence: Confidence score for this fact (0.0–1.0).
    """

    fact_text: str
    subject: str
    action: Literal["ADD", "UPDATE", "NOOP", "DELETE"]
    matched_fact_id: uuid_mod.UUID | None = None
    matched_fact_text: str | None = None
    confidence: float = 0.8


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
@dataclass
class _FactWithSimilarity:
    """New fact paired with its most similar existing facts."""

    extracted_fact: ExtractedFact
    similar_existing: list[tuple[MemoryFact, float]]

    @property
    def max_similarity(self) -> float:
        if not self.similar_existing:
            return 0.0
        return max(score for _, score in self.similar_existing)


async def _fetch_entity_facts(
    session: AsyncSession,
    user_id: str,
    entity_key: str,
) -> list[MemoryFact]:
    """Fetch active facts for an entity."""
    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.entity_key == entity_key,
            MemoryFact.valid_to.is_(None),
        )
        .order_by(MemoryFact.valid_from.desc())
        .limit(50)
    )
    result = await session.execute(stmt)
    return list(result.scalars().all())


async def _find_similar_facts(
    session: AsyncSession,
    new_facts: list[ExtractedFact],
    existing_facts: list[MemoryFact],
    user_id: str,
    entity_key: str,
    embeddings: EmbeddingProvider,
) -> list[_FactWithSimilarity]:
    """For each new fact, find top-3 similar existing facts by embedding."""
    results = []
    for new_fact in new_facts:
        fact_text = f"{new_fact.subject}: {new_fact.fact}"
        embedding = await embeddings.embed_one(fact_text)

        similar: list[tuple[MemoryFact, float]] = []

        if embedding is not None and existing_facts:
            query = text("""
                SELECT id,
                       1 - (embedding_vec <=> CAST(:embedding AS vector)) as similarity
                FROM memory_facts
                WHERE user_id = :user_id
                  AND entity_key = :entity_key
                  AND valid_to IS NULL
                  AND embedding_vec IS NOT NULL
                ORDER BY embedding_vec <=> CAST(:embedding AS vector)
                LIMIT 3
            """)
            try:
                async with session.begin_nested():
                    result = await session.execute(
                        query,
                        {
                            "user_id": user_id,
                            "entity_key": entity_key,
                            "embedding": str(embedding),
                        },
                    )
                    rows = result.fetchall()
            except Exception as e:
                logger.warning("reconcile: pgvector similarity query failed: %s", e)
                rows = []

            existing_by_id = {str(f.id): f for f in existing_facts}
            for row_id, similarity in rows:
                fact_obj = existing_by_id.get(str(row_id))
                if fact_obj:
                    similar.append((fact_obj, float(similarity)))

        results.append(
            _FactWithSimilarity(
                extracted_fact=new_fact,
                similar_existing=similar,
            )
        )

    return results


# ---------------------------------------------------------------------------
# LLM reconciliation
# ---------------------------------------------------------------------------
RECONCILIATION_SYSTEM_PROMPT = """You are a memory reconciliation engine.
Given new facts and existing facts about the same subject, decide what to do.

For each new fact, decide:
- ADD: genuinely new information, not covered by existing facts
- UPDATE: replaces/updates an existing fact (specify which one by number)
- NOOP: already known (same information as an existing fact, specify which)

Respond as JSON:
{
  "decisions": [
    {"new_fact_index": 1, "action": "ADD"},
    {"new_fact_index": 2, "action": "UPDATE", "existing_fact_index": 1},
    {"new_fact_index": 3, "action": "NOOP", "existing_fact_index": 2}
  ]
}"""


async def _llm_reconcile(
    subject_entity: ResolvedEntity,
    fact_candidates: list[_FactWithSimilarity],
    llm: LLMProvider,
) -> list[ReconciliationDecision]:
    """Ask LLM to decide ADD/UPDATE/NOOP for facts with similar existing ones.

    One call per subject (batched). Fallback: all ADD.
    """
    new_facts_lines = []
    for i, fc in enumerate(fact_candidates, 1):
        new_facts_lines.append(f'{i}. "{fc.extracted_fact.fact}" (confidence: {fc.extracted_fact.confidence})')

    all_existing: dict[str, tuple[MemoryFact, int]] = {}
    idx = 1
    for fc in fact_candidates:
        for existing_fact, _ in fc.similar_existing:
            fid = str(existing_fact.id)
            if fid not in all_existing:
                all_existing[fid] = (existing_fact, idx)
                idx += 1

    existing_facts_lines = []
    for _fid, (fact, display_idx) in all_existing.items():
        fact_display = fact.fact_text if fact.fact_text else f"{fact.attribute_key}: {fact.value_text}"
        existing_facts_lines.append(f'{display_idx}. "{fact_display}" (confidence: {fact.confidence})')

    user_prompt = (
        f"Subject: {subject_entity.display_name} ({subject_entity.entity_type})\n\n"
        f"New facts:\n" + "\n".join(new_facts_lines) + "\n\n"
        "Existing facts in memory:\n" + "\n".join(existing_facts_lines)
    )

    try:
        raw = await llm.complete(
            messages=[
                {"role": "system", "content": RECONCILIATION_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            response_format={"type": "json_object"},
            temperature=0,
        )
    except Exception as e:
        logger.warning("reconcile: LLM call failed: %s", e)
        return [
            ReconciliationDecision(
                fact_text=fc.extracted_fact.fact,
                subject=fc.extracted_fact.subject,
                action="ADD",
                confidence=fc.extracted_fact.confidence,
            )
            for fc in fact_candidates
        ]

    # Parse response
    try:
        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)
        llm_decisions = data.get("decisions", [])
    except Exception as e:
        logger.warning("reconcile: LLM response parse failed: %s", e)
        return [
            ReconciliationDecision(
                fact_text=fc.extracted_fact.fact,
                subject=fc.extracted_fact.subject,
                action="ADD",
                confidence=fc.extracted_fact.confidence,
            )
            for fc in fact_candidates
        ]

    existing_by_display_idx = {idx: fact for _, (fact, idx) in all_existing.items()}
    decisions = []

    for d in llm_decisions:
        new_idx = d.get("new_fact_index", 0)
        if not (1 <= new_idx <= len(fact_candidates)):
            continue

        fc = fact_candidates[new_idx - 1]
        action = d.get("action", "ADD").upper()
        if action not in ("ADD", "UPDATE", "NOOP", "DELETE"):
            action = "ADD"

        matched_id = None
        matched_text = None
        existing_idx = d.get("existing_fact_index")
        if existing_idx and existing_idx in existing_by_display_idx:
            matched_fact = existing_by_display_idx[existing_idx]
            matched_id = matched_fact.id
            matched_text = (
                matched_fact.fact_text
                if matched_fact.fact_text
                else f"{matched_fact.attribute_key}: {matched_fact.value_text}"
            )

        decisions.append(
            ReconciliationDecision(
                fact_text=fc.extracted_fact.fact,
                subject=fc.extracted_fact.subject,
                action=action,
                matched_fact_id=matched_id,
                matched_fact_text=matched_text,
                confidence=fc.extracted_fact.confidence,
            )
        )

    # Handle uncovered new facts (default ADD)
    covered_indices = {d.get("new_fact_index", 0) for d in llm_decisions}
    for i, fc in enumerate(fact_candidates, 1):
        if i not in covered_indices:
            decisions.append(
                ReconciliationDecision(
                    fact_text=fc.extracted_fact.fact,
                    subject=fc.extracted_fact.subject,
                    action="ADD",
                    confidence=fc.extracted_fact.confidence,
                )
            )

    return decisions


# ---------------------------------------------------------------------------
# Subject batch reconciliation
# ---------------------------------------------------------------------------
async def _reconcile_subject_batch(
    session: AsyncSession,
    user_id: str,
    subject_entity: ResolvedEntity,
    new_facts: list[ExtractedFact],
    llm: LLMProvider,
    embeddings: EmbeddingProvider,
) -> list[ReconciliationDecision]:
    """Reconcile a batch of facts for a single subject entity."""
    existing = await _fetch_entity_facts(session, user_id, subject_entity.entity_key)

    # No existing facts → all ADD (skip LLM)
    if not existing:
        return [
            ReconciliationDecision(
                fact_text=f.fact,
                subject=f.subject,
                action="ADD",
                confidence=f.confidence,
            )
            for f in new_facts
        ]

    # Find similar existing facts for each new fact
    candidates = await _find_similar_facts(
        session,
        new_facts,
        existing,
        user_id,
        subject_entity.entity_key,
        embeddings,
    )

    # Split: needs LLM vs auto ADD
    needs_llm = [c for c in candidates if c.max_similarity >= SIMILARITY_THRESHOLD]
    no_llm = [c for c in candidates if c.max_similarity < SIMILARITY_THRESHOLD]

    decisions: list[ReconciliationDecision] = [
        ReconciliationDecision(
            fact_text=c.extracted_fact.fact,
            subject=c.extracted_fact.subject,
            action="ADD",
            confidence=c.extracted_fact.confidence,
        )
        for c in no_llm
    ]

    if needs_llm:
        llm_decisions = await _llm_reconcile(subject_entity, needs_llm, llm)
        decisions.extend(llm_decisions)

    return decisions


# ---------------------------------------------------------------------------
# Main entry point
# ---------------------------------------------------------------------------
async def reconcile_facts(
    session: AsyncSession,
    user_id: str,
    new_facts: list[ExtractedFact],
    entity_map: dict[str, ResolvedEntity],
    llm: LLMProvider,
    embeddings: EmbeddingProvider,
) -> list[ReconciliationDecision]:
    """Compare new facts against existing ones and decide actions.

    Groups facts by subject and makes one LLM call per subject.
    If LLM fails, defaults to ADD (prefer duplicates over lost data).

    Args:
        session: Database session.
        user_id: User identifier.
        new_facts: Facts from extraction.
        entity_map: Resolved entities mapping.
        llm: Injected LLM provider.
        embeddings: Injected embedding provider.

    Returns:
        List of ReconciliationDecision, one per input fact.
    """
    if not new_facts:
        return []

    by_subject: dict[str, list[ExtractedFact]] = defaultdict(list)
    for fact in new_facts:
        subj_key = fact.subject.lower().strip()
        by_subject[subj_key].append(fact)

    all_decisions: list[ReconciliationDecision] = []

    for subj_key, facts in by_subject.items():
        entity = entity_map.get(subj_key)
        if entity is None:
            logger.warning("reconcile: no entity for subject '%s', defaulting to ADD", subj_key)
            for f in facts:
                all_decisions.append(
                    ReconciliationDecision(
                        fact_text=f.fact,
                        subject=f.subject,
                        action="ADD",
                        confidence=f.confidence,
                    )
                )
            continue

        try:
            async with session.begin_nested():
                decisions = await _reconcile_subject_batch(
                    session,
                    user_id,
                    entity,
                    facts,
                    llm,
                    embeddings,
                )
            all_decisions.extend(decisions)
        except Exception as e:
            logger.error("reconcile: failed for subject '%s': %s", subj_key, e)
            for f in facts:
                all_decisions.append(
                    ReconciliationDecision(
                        fact_text=f.fact,
                        subject=f.subject,
                        action="ADD",
                        confidence=f.confidence,
                    )
                )

    logger.info(
        "reconcile: user_id=%s total=%d ADD=%d UPDATE=%d NOOP=%d DELETE=%d",
        user_id,
        len(all_decisions),
        sum(1 for d in all_decisions if d.action == "ADD"),
        sum(1 for d in all_decisions if d.action == "UPDATE"),
        sum(1 for d in all_decisions if d.action == "NOOP"),
        sum(1 for d in all_decisions if d.action == "DELETE"),
    )

    return all_decisions
