"""Memify pipeline — vitality scoring, staleness marking, edge management.

Self-improving memory pipeline that computes vitality scores, detects stale
facts, reinforces/decays KG edge strengths, and detects entity merges.
"""

from __future__ import annotations

import asyncio
import json
import logging
import math
from dataclasses import dataclass
from datetime import UTC, datetime, timedelta
from typing import Any

from sqlalchemy import select, update
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.orm import defer

from arandu._helpers import cosine_similarity
from arandu.config import MemoryConfig
from arandu.models import MemoryEntityRelationship, MemoryFact
from arandu.protocols import EmbeddingProvider, LLMProvider

logger = logging.getLogger(__name__)

__all__ = [
    "MemifyResult",
    "compute_vitality",
    "run_memify",
]

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

EDGE_REINFORCE_AMOUNT = 0.05
EDGE_DECAY_FACTOR = 0.9
EDGE_MIN_STRENGTH = 0.1
EDGE_UNUSED_DAYS = 90
EDGE_USED_DAYS = 7
LLM_MERGE_TIMEOUT_SEC = 5.0

DEFAULT_VITALITY_WEIGHTS: dict[str, float] = {
    "retrieval": 0.30,
    "recency": 0.25,
    "confidence": 0.20,
    "reinforcement": 0.15,
}


# ---------------------------------------------------------------------------
# Result dataclass
# ---------------------------------------------------------------------------


@dataclass
class MemifyResult:
    """Result of the memify pipeline.

    Attributes:
        facts_scored: Number of facts scored for vitality.
        facts_marked_stale: Number of facts marked as stale.
        edges_reinforced: Number of KG edges reinforced.
        merges_executed: Number of entity merges executed.
    """

    facts_scored: int = 0
    facts_marked_stale: int = 0
    edges_reinforced: int = 0
    merges_executed: int = 0


# ---------------------------------------------------------------------------
# Vitality scoring (pure function)
# ---------------------------------------------------------------------------


def compute_vitality(
    fact: object,
    now: datetime | None = None,
    *,
    weights: dict[str, float] | None = None,
) -> float:
    """Compute vitality score for a single fact (0.0-1.0).

    Components:
    - retrieval (0.30): log-scale of times_retrieved
    - recency (0.25): exponential decay from last_retrieved_at
    - confidence (0.20): raw confidence value
    - reinforcement (0.15): bounded reinforcement count

    Correction penalty: 0.8 ^ user_correction_count.
    Superseded facts (valid_to is not None) return 0.0.

    Args:
        fact: A MemoryFact-like object.
        now: Current timestamp (defaults to UTC now).
        weights: Optional custom weight dict.

    Returns:
        Vitality score between 0.0 and 1.0.
    """
    if now is None:
        now = datetime.now(UTC)

    # Superseded facts get 0
    valid_to = getattr(fact, "valid_to", None)
    if valid_to is not None:
        return 0.0

    w = weights or DEFAULT_VITALITY_WEIGHTS

    # Retrieval: log-scale normalization
    times_retrieved = getattr(fact, "times_retrieved", 0) or 0
    retrieval_score = math.log(1 + times_retrieved) / math.log(11) if times_retrieved > 0 else 0.0
    retrieval_score = min(1.0, retrieval_score)

    # Recency: exponential decay
    last_retrieved = getattr(fact, "last_retrieved_at", None)
    if last_retrieved is not None:
        if last_retrieved.tzinfo is None:
            last_retrieved = last_retrieved.replace(tzinfo=UTC)
        days = max(0, (now - last_retrieved).total_seconds() / 86400)
        recency_score = math.exp(-days * math.log(2) / 30)
    else:
        # Never retrieved — use creation age with 90-day halflife
        created = getattr(fact, "created_at", None) or getattr(fact, "valid_from", None)
        if created is not None:
            if created.tzinfo is None:
                created = created.replace(tzinfo=UTC)
            days = max(0, (now - created).total_seconds() / 86400)
            recency_score = math.exp(-days * math.log(2) / 90)
        else:
            recency_score = 0.5

    # Confidence
    confidence = float(getattr(fact, "confidence", 0.5) or 0.5)
    confidence_score = max(0.0, min(1.0, confidence))

    # Reinforcement
    reinforcement_count = getattr(fact, "reinforcement_count", 0) or 0
    reinforcement_score = min(1.0, reinforcement_count / 5.0) if reinforcement_count > 0 else 0.0

    # Weighted sum
    raw = (
        w.get("retrieval", 0.30) * retrieval_score
        + w.get("recency", 0.25) * recency_score
        + w.get("confidence", 0.20) * confidence_score
        + w.get("reinforcement", 0.15) * reinforcement_score
    )

    # Correction penalty
    corrections = getattr(fact, "user_correction_count", 0) or 0
    if corrections > 0:
        raw *= 0.8**corrections

    return max(0.0, min(1.0, raw))


# ---------------------------------------------------------------------------
# Protection check
# ---------------------------------------------------------------------------


def _is_fact_protected(fact: object) -> bool:
    """Check if a fact should never be marked stale."""
    importance = getattr(fact, "importance", 0.0) or 0.0
    is_sensitive = getattr(fact, "is_sensitive", False)
    return importance >= 0.8 or bool(is_sensitive)


# ---------------------------------------------------------------------------
# Main memify pipeline
# ---------------------------------------------------------------------------


async def run_memify(
    session: AsyncSession,
    user_id: str,
    llm_provider: LLMProvider,
    embedding_provider: EmbeddingProvider,
    config: MemoryConfig,
) -> MemifyResult:
    """Run the memify pipeline: vitality scoring → staleness → edge management.

    Step 1: Score all active facts for vitality; mark stale if below threshold.
    Step 2: Reinforce/decay KG edge strengths based on recency of use.
    Step 3: (Optional) Detect and execute entity merges.

    Args:
        session: Database session.
        user_id: User identifier.
        llm_provider: Injected LLM provider.
        embedding_provider: Injected embedding provider.
        config: Memory configuration.

    Returns:
        MemifyResult with pipeline stats.
    """
    now = datetime.now(UTC)
    stats = MemifyResult()

    # Step 1: Vitality scoring
    facts_stmt = (
        select(MemoryFact)
        .where(MemoryFact.user_id == user_id, MemoryFact.valid_to.is_(None))
        .options(defer(MemoryFact.embedding))
    )
    facts_result = await session.execute(facts_stmt)
    facts = list(facts_result.scalars().all())

    for fact in facts:
        vitality = compute_vitality(fact, now)
        stats.facts_scored += 1

        if vitality < config.vitality_stale_threshold and not _is_fact_protected(fact):
            fact.is_stale = True
            stats.facts_marked_stale += 1

    # Step 2: Edge reinforcement/decay
    edges_stmt = select(MemoryEntityRelationship).where(MemoryEntityRelationship.user_id == user_id)
    edges_result = await session.execute(edges_stmt)
    edges = list(edges_result.scalars().all())

    cutoff_used = now - timedelta(days=EDGE_USED_DAYS)
    cutoff_unused = now - timedelta(days=EDGE_UNUSED_DAYS)

    for edge in edges:
        last_used = getattr(edge, "last_used_at", None)
        if last_used is not None:
            if last_used.tzinfo is None:
                last_used = last_used.replace(tzinfo=UTC)
            if last_used >= cutoff_used:
                edge.strength = min(1.0, (edge.strength or 0.5) + EDGE_REINFORCE_AMOUNT)
                stats.edges_reinforced += 1
            elif last_used < cutoff_unused:
                edge.strength = max(EDGE_MIN_STRENGTH, (edge.strength or 0.5) * EDGE_DECAY_FACTOR)

    # Step 3: Entity merge detection (via embedding similarity + LLM confirmation)
    merges = await _detect_and_merge_entities(
        session,
        user_id,
        facts,
        llm_provider,
        embedding_provider,
        config,
        now,
    )
    stats.merges_executed = merges

    logger.info(
        "run_memify: user=%s, scored=%d, stale=%d, edges_reinforced=%d, merges=%d",
        user_id,
        stats.facts_scored,
        stats.facts_marked_stale,
        stats.edges_reinforced,
        stats.merges_executed,
    )
    return stats


# ---------------------------------------------------------------------------
# Entity merge detection
# ---------------------------------------------------------------------------


async def _detect_and_merge_entities(
    session: AsyncSession,
    user_id: str,
    active_facts: list[Any],
    llm_provider: LLMProvider,
    embedding_provider: EmbeddingProvider,
    config: MemoryConfig,
    now: datetime,
) -> int:
    """Detect and merge duplicate entities via embedding similarity + LLM."""
    # Build entity embeddings by averaging fact embeddings
    entity_embeddings: dict[str, list[list[float]]] = {}
    entity_types: dict[str, str] = {}
    for fact in active_facts:
        key = fact.entity_key
        etype = fact.entity_type
        emb = getattr(fact, "embedding", None)
        if key and emb and isinstance(emb, list) and len(emb) > 0:
            entity_embeddings.setdefault(key, []).append(emb)
            entity_types[key] = etype or ""

    # Average embeddings per entity
    avg_embeddings: dict[str, list[float]] = {}
    for key, embs in entity_embeddings.items():
        if not embs:
            continue
        dim = len(embs[0])
        avg = [sum(e[d] for e in embs) / len(embs) for d in range(dim)]
        avg_embeddings[key] = avg

    # Find merge candidates (same type, high similarity)
    keys = list(avg_embeddings.keys())
    merges_executed = 0

    for i in range(len(keys)):
        for j in range(i + 1, len(keys)):
            k1, k2 = keys[i], keys[j]
            if entity_types.get(k1) != entity_types.get(k2):
                continue
            sim = cosine_similarity(avg_embeddings[k1], avg_embeddings[k2])
            if sim < config.memify_merge_similarity_threshold:
                continue

            # LLM confirmation
            confirmed = await _llm_confirm_merge(
                session,
                user_id,
                k1,
                k2,
                sim,
                active_facts,
                llm_provider,
            )
            if confirmed:
                await _execute_merge(session, user_id, k1, k2)
                merges_executed += 1

    return merges_executed


async def _llm_confirm_merge(
    session: AsyncSession,
    user_id: str,
    ek1: str,
    ek2: str,
    similarity: float,
    active_facts: list[Any],
    llm: LLMProvider,
) -> bool:
    """Confirm entity merge via LLM."""
    facts1 = [f for f in active_facts if f.entity_key == ek1][:10]
    facts2 = [f for f in active_facts if f.entity_key == ek2][:10]

    blob1 = "\n".join(f"- {f.attribute_key}: {f.value_text or ''}" for f in facts1)
    blob2 = "\n".join(f"- {f.attribute_key}: {f.value_text or ''}" for f in facts2)

    messages = [
        {
            "role": "system",
            "content": (
                "Analyze if these two entities are the SAME person/thing. "
                'Return JSON: {"should_merge": true/false, "reason": "brief explanation"}. '
                "ONLY JSON."
            ),
        },
        {
            "role": "user",
            "content": (f"Entity 1: {ek1}\n{blob1}\n\nEntity 2: {ek2}\n{blob2}\n\nSimilarity: {similarity:.3f}"),
        },
    ]
    try:
        _llm_result = await asyncio.wait_for(llm.complete(messages), timeout=LLM_MERGE_TIMEOUT_SEC)
        raw = _llm_result.text
        if not raw:
            return False
        raw = raw.strip()
        if raw.startswith("```"):
            raw = raw.split("\n", 1)[-1] if "\n" in raw else raw[3:]
        if raw.endswith("```"):
            raw = raw.rsplit("```", 1)[0].strip()
        data = json.loads(raw)
        return bool(data.get("should_merge", False))
    except Exception as e:
        logger.debug("llm_confirm_merge_failed: %s", e)
        return False


async def _execute_merge(
    session: AsyncSession,
    user_id: str,
    winner_key: str,
    loser_key: str,
) -> None:
    """Execute entity merge: reassign facts and relationships."""
    # Update facts
    await session.execute(
        update(MemoryFact)
        .where(MemoryFact.user_id == user_id, MemoryFact.entity_key == loser_key)
        .values(entity_key=winner_key)
    )

    # Update relationships (source)
    await session.execute(
        update(MemoryEntityRelationship)
        .where(
            MemoryEntityRelationship.user_id == user_id,
            MemoryEntityRelationship.source_entity_key == loser_key,
        )
        .values(source_entity_key=winner_key)
    )

    # Update relationships (target)
    await session.execute(
        update(MemoryEntityRelationship)
        .where(
            MemoryEntityRelationship.user_id == user_id,
            MemoryEntityRelationship.target_entity_key == loser_key,
        )
        .values(target_entity_key=winner_key)
    )

    logger.info("entity_merge_executed: user=%s, winner=%s, loser=%s", user_id, winner_key, loser_key)
