"""Write pipeline orchestrator — extract → resolve → reconcile → upsert.

Ported from advisor_api memory_service.py. Simplified: no trace, no runtime config.
All dependencies injected via function parameters.
"""

import json
import logging
import time
import traceback
from dataclasses import dataclass, field
from datetime import UTC, datetime

from sqlalchemy.ext.asyncio import AsyncSession

from arandu.config import MemoryConfig
from arandu.constants import ExtractedFact, ExtractionOutput
from arandu.models import MemoryEvent
from arandu.protocols import EmbeddingProvider, LLMProvider
from arandu.tracing import PipelineTrace
from arandu.write.entity_helpers import create_or_update_entity
from arandu.write.entity_resolution import resolve_entities
from arandu.write.extract import _call_llm, _strip_markdown_fences, extract_from_message
from arandu.write.reconcile import reconcile_facts
from arandu.write.upsert import UpsertResult, execute_upsert, upsert_relations

logger = logging.getLogger(__name__)

EMOTION_DETECTION_PROMPT = """You are an emotion classifier for personal messages.
Given a user message, classify the primary emotion and intensity.

Respond in JSON with keys:
- emotion_primary: one of "joy", "sadness", "anger", "fear", "surprise", "disgust", "trust", "anticipation", "neutral"
- emotion_intensity: float 0.0 to 1.0 (0.0 = barely perceptible, 1.0 = overwhelming)
- energy_level: one of "high", "medium", "low"

If the message has no clear emotion, use "neutral" with intensity 0.0 and energy "medium"."""


async def _detect_emotion(
    message: str,
    llm: LLMProvider,
    event: MemoryEvent,
) -> None:
    """Detect emotion from message and update event fields. Fails silently."""
    try:
        raw = await _call_llm(
            llm=llm,
            system_prompt=EMOTION_DETECTION_PROMPT,
            user_prompt=message,
            call_type="emotion_detection",
        )
        if raw is None:
            return

        cleaned = _strip_markdown_fences(raw)
        data = json.loads(cleaned)

        event.emotion_primary = data.get("emotion_primary")
        event.emotion_intensity = data.get("emotion_intensity")
        event.energy_level = data.get("energy_level")
    except Exception as e:
        logger.debug("write_pipeline: emotion detection failed (non-critical): %s", e)


@dataclass
class WritePipelineResult:
    """Internal result of the write pipeline. Mapped to WriteResult by MemoryClient."""

    event_id: str = ""
    facts_added: list = field(default_factory=list)
    facts_updated: list = field(default_factory=list)
    facts_unchanged: list = field(default_factory=list)
    facts_deleted: list = field(default_factory=list)
    entities_resolved: list = field(default_factory=list)
    duration_ms: float = 0.0
    success: bool = True
    error: str | None = None


async def _semantic_dedup_facts(
    facts: list[ExtractedFact],
    embeddings: EmbeddingProvider,
    threshold: float = 0.85,
) -> list[ExtractedFact]:
    """Remove semantically duplicate facts via embedding cosine similarity.

    Compares all fact texts pairwise. When two facts have similarity >= threshold,
    keeps the first one (which tends to be from the primary subject since extraction
    processes entities in order). Fail-safe: returns original list on any error.
    """
    if len(facts) <= 1:
        return facts

    try:
        texts = [f.fact for f in facts]
        vecs = await embeddings.embed(texts)
        if not vecs or len(vecs) != len(facts):
            return facts

        # Cosine similarity between all pairs
        import math

        def _cosine(a: list[float], b: list[float]) -> float:
            dot = sum(x * y for x, y in zip(a, b))
            na = math.sqrt(sum(x * x for x in a))
            nb = math.sqrt(sum(x * x for x in b))
            if na == 0 or nb == 0:
                return 0.0
            return dot / (na * nb)

        # Mark duplicates (keep first occurrence)
        removed: set[int] = set()
        for i in range(len(facts)):
            if i in removed:
                continue
            for j in range(i + 1, len(facts)):
                if j in removed:
                    continue
                sim = _cosine(vecs[i], vecs[j])
                if sim >= threshold:
                    removed.add(j)

        deduped = [f for idx, f in enumerate(facts) if idx not in removed]
        if len(deduped) < len(facts):
            logger.info(
                "semantic_dedup: removed %d/%d duplicate facts (threshold=%.2f)",
                len(facts) - len(deduped),
                len(facts),
                threshold,
            )
        return deduped

    except Exception as e:
        logger.warning("semantic_dedup: failed, returning original facts: %s", e)
        return facts


async def _create_event(
    session: AsyncSession,
    user_id: str,
    text: str,
    source: str,
    embeddings: EmbeddingProvider,
) -> MemoryEvent:
    """Create an immutable event record for a user message."""
    embedding = await embeddings.embed_one(text)

    event = MemoryEvent(
        user_id=user_id,
        occurred_at=datetime.now(UTC),
        text=text,
        source=source,
        importance=0.5,
        embedding_vec=embedding,
    )
    session.add(event)
    await session.flush()
    return event


async def run_write_pipeline(
    session: AsyncSession,
    user_id: str,
    message: str,
    llm: LLMProvider,
    embeddings: EmbeddingProvider,
    config: MemoryConfig,
    source: str = "api",
    recent_messages: list[str] | None = None,
    trace: PipelineTrace | None = None,
    dry_run: bool = False,
) -> WritePipelineResult:
    """Execute the full write pipeline: extract → resolve → reconcile → upsert.

    Args:
        session: Database session (caller manages transaction/commit).
        user_id: Unique identifier for the user.
        message: The user's message text.
        llm: Injected LLM provider.
        embeddings: Injected embedding provider.
        config: Memory configuration.
        source: Source channel identifier.
        recent_messages: Optional conversation context (last N messages)
            for resolving pronouns and anaphora references in extraction.

    Returns:
        Dict with pipeline results (used by MemoryClient.write to build WriteResult).
    """
    t_start = time.perf_counter()

    # 0. Guard: empty/whitespace messages return immediately
    if not message or not message.strip():
        duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
        logger.info("write_pipeline: empty message, skipping (user=%s)", user_id)
        return WritePipelineResult(duration_ms=duration_ms)

    # 1. Create immutable event (skip in dry_run — zero side effects)
    event_id = ""
    if not dry_run:
        event = await _create_event(session, user_id, message, source, embeddings)
        event_id = str(event.id)
        logger.info("write_pipeline: user=%s, event_id=%s", user_id, event_id)

        # 1b. Detect emotion (non-critical, fails silently)
        await _detect_emotion(message, llm, event)
    else:
        logger.info("write_pipeline: dry_run=True, skipping event creation (user=%s)", user_id)

    # Wrap pipeline steps 2-6 in error handling — event is already persisted
    # and survives even if the pipeline fails.
    extraction = ExtractionOutput()
    upsert_result = UpsertResult()

    try:
        # 2. Extract entities, facts, relations (outside savepoint — no DB writes)
        _t_extract = time.perf_counter() if trace else 0.0
        extraction, extraction_meta = await extract_from_message(
            message,
            llm=llm,
            config=config,
            recent_messages=recent_messages,
        )
        if trace:
            trace.add_step(
                "extraction",
                {
                    "relation_retry_triggered": extraction_meta.get("relation_retry_triggered", False),
                    "entities": [e.model_dump() for e in extraction.entities],
                    "facts": [f.model_dump() for f in extraction.facts],
                    "relations": [r.model_dump() for r in extraction.relations],
                    "entity_count": len(extraction.entities),
                    "fact_count": len(extraction.facts),
                    "relation_count": len(extraction.relations),
                },
                trace.elapsed(_t_extract),
            )

        # 2b. Semantic dedup: remove near-duplicate facts by embedding similarity
        if len(extraction.facts) > 1:
            original_count = len(extraction.facts)
            extraction.facts = await _semantic_dedup_facts(
                extraction.facts,
                embeddings,
                threshold=0.85,
            )
            if trace and len(extraction.facts) < original_count:
                trace.add_step(
                    "semantic_dedup",
                    {
                        "before": original_count,
                        "after": len(extraction.facts),
                        "removed": original_count - len(extraction.facts),
                    },
                    0.0,
                )

        if not extraction.facts and not extraction.relations:
            duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
            return WritePipelineResult(event_id=event_id, duration_ms=duration_ms)

        # dry_run: return extraction results without persisting
        if dry_run:
            duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
            return WritePipelineResult(
                event_id="",
                facts_added=[f.fact for f in extraction.facts],
                entities_resolved=[e.name for e in extraction.entities],
                duration_ms=duration_ms,
            )

        # Steps 3-6 run inside a savepoint so that if anything fails,
        # the partial DB writes are rolled back cleanly but the event
        # (flushed above) survives. Note: inner functions (entity_resolution,
        # reconcile, upsert) also use their own savepoints for per-operation
        # resilience. This outer savepoint provides pipeline-level atomicity.
        async with session.begin_nested():
            # 3. Resolve entities
            _t_er = time.perf_counter() if trace else 0.0
            entity_map = await resolve_entities(
                session,
                user_id,
                extraction.entities,
                extraction.facts,
                llm=llm,
                embeddings=embeddings,
                fuzzy_threshold=config.fuzzy_threshold,
                enable_llm_resolution=config.enable_llm_resolution,
            )
            if trace:
                trace.add_step(
                    "entity_resolution",
                    {
                        "entities": [
                            {
                                "name": e.display_name,
                                "key": e.entity_key,
                                "method": e.resolution_method,
                                "is_new": e.is_new,
                                "fuzzy_score": e.fuzzy_score,
                            }
                            for e in entity_map.values()
                        ],
                        "entity_count": len(entity_map),
                    },
                    trace.elapsed(_t_er),
                )

            # 3b. Ensure all resolved entities exist in memory_entities table
            for resolved in entity_map.values():
                try:
                    await create_or_update_entity(
                        session,
                        user_id,
                        resolved.entity_key,
                        display_name=resolved.display_name,
                        entity_type=resolved.entity_type,
                    )
                except Exception as e:
                    logger.warning(
                        "write_pipeline: entity upsert failed for %s (non-critical): %s",
                        resolved.entity_key,
                        e,
                    )

            # 4. Reconcile facts against existing
            _t_rec = time.perf_counter() if trace else 0.0
            decisions = await reconcile_facts(
                session,
                user_id,
                extraction.facts,
                entity_map,
                llm=llm,
                embeddings=embeddings,
            )
            if trace:
                trace.add_step(
                    "reconciliation",
                    {
                        "decisions": [
                            {
                                "fact_text": d.fact_text,
                                "subject": d.subject,
                                "action": d.action,
                                "matched_fact_id": str(d.matched_fact_id) if d.matched_fact_id else None,
                                "matched_fact_text": d.matched_fact_text,
                                "confidence": d.confidence,
                            }
                            for d in decisions
                        ],
                        "action_counts": {
                            "ADD": sum(1 for d in decisions if d.action == "ADD"),
                            "UPDATE": sum(1 for d in decisions if d.action == "UPDATE"),
                            "NOOP": sum(1 for d in decisions if d.action == "NOOP"),
                            "DELETE": sum(1 for d in decisions if d.action == "DELETE"),
                        },
                    },
                    trace.elapsed(_t_rec),
                )

            # 5. Execute decisions (upsert/close/confirm)
            _t_upsert = time.perf_counter() if trace else 0.0
            upsert_result = await execute_upsert(
                session,
                user_id,
                event_id,
                decisions,
                entity_map,
                embeddings=embeddings,
            )
            if trace:
                trace.add_step(
                    "upsert",
                    {
                        "added": upsert_result.added,
                        "updated": upsert_result.updated,
                        "confirmed": len(upsert_result.confirmed),
                        "deleted": len(upsert_result.deleted),
                    },
                    trace.elapsed(_t_upsert),
                )

            # 6. Upsert relationships (with evidence linkage + mirror facts)
            if extraction.relations:
                await upsert_relations(
                    session,
                    user_id,
                    event_id,
                    extraction.relations,
                    entity_map,
                    created_facts=upsert_result.created_facts,
                    embeddings=embeddings,
                )

    except Exception as e:
        logger.warning(
            "write_pipeline: failed after event creation: user=%s event=%s error=%s",
            user_id,
            event_id,
            e,
            exc_info=True,
        )
        if trace:
            trace.add_step(
                "error",
                {
                    "error": str(e),
                    "error_type": type(e).__name__,
                    "traceback": traceback.format_exc(),
                },
                0.0,
            )
        duration_ms = round((time.perf_counter() - t_start) * 1000, 1)
        return WritePipelineResult(
            event_id=event_id,
            duration_ms=duration_ms,
            success=False,
            error=f"{type(e).__name__}: {e}",
        )

    duration_ms = round((time.perf_counter() - t_start) * 1000, 1)

    logger.info(
        "write_pipeline: user=%s, event=%s, added=%d, updated=%d, confirmed=%d, entities=%d, latency_ms=%.1f",
        user_id,
        event_id,
        upsert_result.added,
        upsert_result.updated,
        len(upsert_result.confirmed),
        len(entity_map),
        duration_ms,
    )

    # Build result data
    facts_added = [
        {"fact_text": d.fact_text, "subject": d.subject, "confidence": d.confidence}
        for d in decisions
        if d.action == "ADD"
    ]
    facts_updated = [
        {"fact_text": d.fact_text, "subject": d.subject, "confidence": d.confidence, "replaced": d.matched_fact_text}
        for d in decisions
        if d.action == "UPDATE"
    ]
    entities_resolved = [
        {"name": e.display_name, "key": e.entity_key, "method": e.resolution_method, "is_new": e.is_new}
        for e in entity_map.values()
    ]

    result = WritePipelineResult(
        event_id=event_id,
        facts_added=facts_added,
        facts_updated=facts_updated,
        facts_unchanged=upsert_result.confirmed,
        facts_deleted=upsert_result.deleted,
        entities_resolved=entities_resolved,
        duration_ms=duration_ms,
    )

    # Persist trace_json on the event for observability
    try:
        event.trace_json = {
            "extraction_mode": "multi_pass",
            "entities_extracted": len(extraction.entities),
            "facts_extracted": len(extraction.facts),
            "relations_extracted": len(extraction.relations),
            "entities_resolved": len(entity_map),
            "decisions": {
                "add": upsert_result.added,
                "update": upsert_result.updated,
                "confirmed": len(upsert_result.confirmed),
                "deleted": len(upsert_result.deleted),
            },
            "duration_ms": duration_ms,
        }
    except Exception as e:
        logger.debug("write_pipeline: trace_json persistence failed: %s", e)

    return result
