"""Sleep-time compute — entity importance, summary refresh, community detection.

Pre-computation jobs that run during maintenance to improve retrieval quality.
"""

from __future__ import annotations

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

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

from arandu.config import MemoryConfig
from arandu.models import MemoryEntity, MemoryEntityRelationship, MemoryFact, MemoryFactEntityLink
from arandu.protocols import LLMProvider

logger = logging.getLogger(__name__)

__all__ = [
    "EntityImportanceResult",
    "SummaryRefreshResult",
    "compute_entity_importance",
    "detect_entity_communities",
    "refresh_entity_summaries",
]

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

IMPORTANCE_WEIGHTS = {
    "fact_density": 0.30,
    "recency": 0.25,
    "retrieval_frequency": 0.25,
    "relationship_degree": 0.20,
}
MAX_ENTITIES_PER_REFRESH = 10
ENTITY_SUMMARY_TIMEOUT_SEC = 5.0


# ---------------------------------------------------------------------------
# Result dataclasses
# ---------------------------------------------------------------------------


@dataclass
class EntityImportanceResult:
    """Result of entity importance scoring.

    Attributes:
        entities_scored: Number of entities scored.
        top_entities: Top entities by score (key, score) pairs.
    """

    entities_scored: int = 0
    top_entities: list[tuple[str, float]] = None  # type: ignore[assignment]

    def __post_init__(self) -> None:
        if self.top_entities is None:
            self.top_entities = []


@dataclass
class SummaryRefreshResult:
    """Result of entity summary refresh.

    Attributes:
        summaries_refreshed: Number of summaries generated.
        summaries_skipped: Number of entities skipped (no facts or LLM failure).
    """

    summaries_refreshed: int = 0
    summaries_skipped: int = 0


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _recency_score(last_seen_at: datetime | None, now: datetime, halflife_days: int = 30) -> float:
    """Exponential decay recency score (0.0-1.0)."""
    if last_seen_at is None:
        return 0.0
    if last_seen_at.tzinfo is None:
        last_seen_at = last_seen_at.replace(tzinfo=UTC)
    days = max(0, (now - last_seen_at).total_seconds() / 86400)
    return math.exp(-days * math.log(2) / max(1, halflife_days))


def _normalize(values: list[float]) -> list[float]:
    """Min-max normalization to [0.0, 1.0]."""
    if not values:
        return []
    mn, mx = min(values), max(values)
    if mx == mn:
        return [0.5] * len(values)
    return [(v - mn) / (mx - mn) for v in values]


# ---------------------------------------------------------------------------
# Entity importance
# ---------------------------------------------------------------------------


async def compute_entity_importance(
    session: AsyncSession,
    user_id: str,
    config: MemoryConfig,
) -> EntityImportanceResult:
    """Compute importance scores for all active entities.

    Formula: weighted sum of fact_density, recency, retrieval_frequency,
    relationship_degree. Updates ``MemoryEntity.importance_score`` in-place.

    Args:
        session: Database session.
        user_id: User identifier.
        config: Memory configuration.

    Returns:
        EntityImportanceResult with stats and top entities.
    """
    now = datetime.now(UTC)

    # Load active entities
    stmt = select(MemoryEntity).where(MemoryEntity.user_id == user_id, MemoryEntity.is_active.is_(True))
    result = await session.execute(stmt)
    entities = list(result.scalars().all())

    if not entities:
        return EntityImportanceResult()

    # Fact stats: count and total retrieved per entity (via links for cross-entity coverage)
    fact_stats_stmt = (
        select(
            MemoryFactEntityLink.entity_key,
            func.count(func.distinct(MemoryFactEntityLink.fact_id)).label("fact_count"),
            func.coalesce(func.sum(MemoryFact.times_retrieved), 0).label("total_retrieved"),
        )
        .join(MemoryFact, MemoryFact.id == MemoryFactEntityLink.fact_id)
        .where(
            MemoryFactEntityLink.user_id == user_id,
            MemoryFact.valid_to.is_(None),
        )
        .group_by(MemoryFactEntityLink.entity_key)
    )
    fact_result = await session.execute(fact_stats_stmt)
    fact_stats: dict[str, tuple[int, int]] = {}
    for row in fact_result.all():
        fact_stats[row[0]] = (row[1], row[2])

    # Fallback: if links table is empty, use direct entity_key
    if not fact_stats:
        fallback_stmt = (
            select(
                MemoryFact.entity_key,
                func.count().label("fact_count"),
                func.coalesce(func.sum(MemoryFact.times_retrieved), 0).label("total_retrieved"),
            )
            .where(MemoryFact.user_id == user_id, MemoryFact.valid_to.is_(None))
            .group_by(MemoryFact.entity_key)
        )
        fallback_result = await session.execute(fallback_stmt)
        for row in fallback_result.all():
            fact_stats[row[0]] = (row[1], row[2])

    # Relationship degree per entity (in + out)
    rel_out_stmt = (
        select(
            MemoryEntityRelationship.source_entity_key,
            func.count().label("cnt"),
        )
        .where(MemoryEntityRelationship.user_id == user_id)
        .group_by(MemoryEntityRelationship.source_entity_key)
    )
    rel_in_stmt = (
        select(
            MemoryEntityRelationship.target_entity_key,
            func.count().label("cnt"),
        )
        .where(MemoryEntityRelationship.user_id == user_id)
        .group_by(MemoryEntityRelationship.target_entity_key)
    )
    rel_out_result = await session.execute(rel_out_stmt)
    rel_in_result = await session.execute(rel_in_stmt)

    rel_degree: dict[str, int] = defaultdict(int)
    for row in rel_out_result.all():  # type: ignore[assignment]
        rel_degree[row[0]] += row[1]
    for row in rel_in_result.all():  # type: ignore[assignment]
        rel_degree[row[0]] += row[1]

    # Compute raw values
    raw_fact_density = [float(fact_stats.get(e.canonical_key, (0, 0))[0]) for e in entities]
    raw_recency = [_recency_score(e.last_seen_at, now, config.importance_recency_halflife_days) for e in entities]
    raw_retrieval = [float(fact_stats.get(e.canonical_key, (0, 0))[1]) for e in entities]
    raw_rel_degree = [float(rel_degree.get(e.canonical_key, 0)) for e in entities]

    # Normalize
    norm_density = _normalize(raw_fact_density)
    norm_recency = _normalize(raw_recency)
    norm_retrieval = _normalize(raw_retrieval)
    norm_rel = _normalize(raw_rel_degree)

    # Weighted sum
    w = IMPORTANCE_WEIGHTS
    scored: list[tuple[str, float]] = []
    for i, entity in enumerate(entities):
        score = (
            w["fact_density"] * norm_density[i]
            + w["recency"] * norm_recency[i]
            + w["retrieval_frequency"] * norm_retrieval[i]
            + w["relationship_degree"] * norm_rel[i]
        )
        score = max(0.0, min(1.0, score))
        entity.importance_score = round(score, 4)
        scored.append((entity.canonical_key, score))

    scored.sort(key=lambda x: x[1], reverse=True)

    logger.info("compute_entity_importance: user=%s, scored=%d", user_id, len(entities))
    return EntityImportanceResult(
        entities_scored=len(entities),
        top_entities=scored[:10],
    )


# ---------------------------------------------------------------------------
# Summary refresh
# ---------------------------------------------------------------------------


async def refresh_entity_summaries(
    session: AsyncSession,
    user_id: str,
    llm_provider: LLMProvider,
    config: MemoryConfig,
) -> SummaryRefreshResult:
    """Refresh entity summaries via LLM for stale entities.

    An entity is stale if its ``summary_text`` is NULL or
    ``summary_refreshed_at`` is older than ``summary_refresh_interval_days``.

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

    Returns:
        SummaryRefreshResult with stats.
    """
    now = datetime.now(UTC)
    cutoff = now - timedelta(days=config.summary_refresh_interval_days)

    # Find stale entities
    stmt = (
        select(MemoryEntity)
        .where(
            MemoryEntity.user_id == user_id,
            MemoryEntity.is_active.is_(True),
        )
        .order_by(MemoryEntity.importance_score.desc().nulls_last())
        .limit(MAX_ENTITIES_PER_REFRESH)
    )
    result = await session.execute(stmt)
    entities = list(result.scalars().all())

    # Filter stale ones
    stale_entities = [
        e
        for e in entities
        if e.summary_text is None
        or getattr(e, "summary_refreshed_at", None) is None
        or (getattr(e, "summary_refreshed_at", None) is not None and e.summary_refreshed_at < cutoff)  # type: ignore[operator]
    ]

    stats = SummaryRefreshResult()

    for entity in stale_entities:
        facts_stmt = (
            select(MemoryFact)
            .where(
                MemoryFact.user_id == user_id,
                MemoryFact.entity_key == entity.canonical_key,
                MemoryFact.valid_to.is_(None),
            )
            .order_by(MemoryFact.valid_from.desc())
            .limit(30)
            .options(defer(MemoryFact.embedding))
        )
        facts_result = await session.execute(facts_stmt)
        facts = list(facts_result.scalars().all())

        if not facts:
            stats.summaries_skipped += 1
            continue

        facts_blob = "\n".join(
            f"- {f.attribute_key}: {f.fact_text or f.value_text or str(f.value_json)}" for f in facts
        )

        messages = [
            {
                "role": "system",
                "content": (
                    "Generate a concise summary (2-3 sentences) of this entity based on the facts. "
                    "Capture who/what it is and the most relevant facts. "
                    "Return ONLY the text, no formatting. "
                    "Write in the same language as the fact content."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Entity: {entity.display_name or entity.canonical_key} "
                    f"(type: {entity.entity_type})\n\nFacts:\n{facts_blob}"
                ),
            },
        ]
        try:
            _llm_result = await asyncio.wait_for(
                llm_provider.complete(messages),
                timeout=ENTITY_SUMMARY_TIMEOUT_SEC,
            )
            if _llm_result.text:
                entity.summary_text = _llm_result.text.strip()[:2000]
                if hasattr(entity, "summary_refreshed_at"):
                    entity.summary_refreshed_at = now
                stats.summaries_refreshed += 1
            else:
                stats.summaries_skipped += 1
        except Exception as e:
            logger.warning("entity_summary_failed: entity=%s error=%s", entity.canonical_key, e)
            stats.summaries_skipped += 1

    logger.info(
        "refresh_entity_summaries: user=%s, refreshed=%d, skipped=%d",
        user_id,
        stats.summaries_refreshed,
        stats.summaries_skipped,
    )
    return stats


# ---------------------------------------------------------------------------
# Community detection (connected-component on entity relationship graph)
# ---------------------------------------------------------------------------


class _UnionFind:
    """Simple Union-Find for grouping entity indices."""

    def __init__(self, n: int) -> None:
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x: int) -> int:
        while self.parent[x] != x:
            self.parent[x] = self.parent[self.parent[x]]
            x = self.parent[x]
        return x

    def union(self, a: int, b: int) -> None:
        ra, rb = self.find(a), self.find(b)
        if ra == rb:
            return
        if self.rank[ra] < self.rank[rb]:
            ra, rb = rb, ra
        self.parent[rb] = ra
        if self.rank[ra] == self.rank[rb]:
            self.rank[ra] += 1

    def groups(self) -> dict[int, list[int]]:
        result: dict[int, list[int]] = defaultdict(list)
        for i in range(len(self.parent)):
            result[self.find(i)].append(i)
        return result


async def detect_entity_communities(
    session: AsyncSession,
    user_id: str,
    config: MemoryConfig,
) -> dict[str, Any]:
    """Connected-component clustering on entity relationship graph.

    Groups entities connected by relationships with strength >= 0.3
    into communities.

    Args:
        session: Database session.
        user_id: User identifier.
        config: Memory configuration.

    Returns:
        Dict with ``communities_found`` and ``total_entities_assigned``.
    """
    # Load entities
    entity_stmt = select(MemoryEntity).where(MemoryEntity.user_id == user_id, MemoryEntity.is_active.is_(True))
    entity_result = await session.execute(entity_stmt)
    entities = list(entity_result.scalars().all())

    if len(entities) < 2:
        return {"communities_found": 0, "total_entities_assigned": 0}

    # Build entity index
    key_to_idx: dict[str, int] = {e.canonical_key: i for i, e in enumerate(entities)}

    # Load relationships
    rel_stmt = select(MemoryEntityRelationship).where(
        MemoryEntityRelationship.user_id == user_id,
        MemoryEntityRelationship.strength >= 0.3,
    )
    rel_result = await session.execute(rel_stmt)
    relationships = list(rel_result.scalars().all())

    if not relationships:
        return {"communities_found": 0, "total_entities_assigned": 0}

    # Union-Find
    n = len(entities)
    uf = _UnionFind(n)
    for rel in relationships:
        src_idx = key_to_idx.get(rel.source_entity_key)
        tgt_idx = key_to_idx.get(rel.target_entity_key)
        if src_idx is not None and tgt_idx is not None:
            uf.union(src_idx, tgt_idx)

    groups = uf.groups()
    communities = [members for members in groups.values() if len(members) >= 2]

    total_assigned = sum(len(c) for c in communities)
    logger.info(
        "detect_entity_communities: user=%s, communities=%d, entities_assigned=%d",
        user_id,
        len(communities),
        total_assigned,
    )
    return {"communities_found": len(communities), "total_entities_assigned": total_assigned}
