"""Memory clustering — semantic grouping of facts for richer retrieval.

Groups facts by ``(entity_type, entity_key)``, creates/updates clusters with
LLM-generated summaries and embeddings. Idempotent. Second-pass community
detection finds cross-entity themes by comparing cluster embeddings.
"""

from __future__ import annotations

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

from sqlalchemy import func, 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 MemoryCluster, MemoryFact, MemoryMetaObservation
from arandu.protocols import EmbeddingProvider, LLMProvider

logger = logging.getLogger(__name__)

__all__ = [
    "ClusteringResult",
    "CommunityDetectionResult",
    "cluster_user_facts",
    "detect_communities",
]

# ---------------------------------------------------------------------------
# Entity label mapping
# ---------------------------------------------------------------------------


def _cluster_label_for_entity(entity_type: str, entity_key: str) -> str:
    """Human-readable label derived from entity_key (no hardcoded map)."""
    combined = f"{entity_type}:{entity_key}" if entity_type and entity_key else entity_key
    if combined == "user:self" or entity_key == "user:self":
        return "User"
    # Extract name part after ":" and title-case it
    if ":" in entity_key:
        name_part = entity_key.split(":", 1)[1]
    else:
        name_part = entity_key or entity_type
    return name_part.replace("_", " ").title() if name_part else entity_type.title()


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


@dataclass
class ClusteringResult:
    """Result of fact clustering.

    Attributes:
        clusters_created: Number of new clusters created.
        clusters_reinforced: Number of existing clusters updated.
        summaries_generated: Number of cluster summaries generated via LLM.
        facts_assigned: Number of facts assigned to clusters.
    """

    clusters_created: int = 0
    clusters_reinforced: int = 0
    summaries_generated: int = 0
    facts_assigned: int = 0


@dataclass
class CommunityDetectionResult:
    """Result of cross-entity community detection.

    Attributes:
        communities_created: New community observations created.
        communities_reinforced: Existing community observations reinforced.
        clusters_in_communities: Total clusters assigned to communities.
        skipped: Whether detection was skipped.
        skip_reason: Reason for skipping, if any.
    """

    communities_created: int = 0
    communities_reinforced: int = 0
    clusters_in_communities: int = 0
    skipped: bool = False
    skip_reason: str | None = None


# ---------------------------------------------------------------------------
# LLM helpers
# ---------------------------------------------------------------------------


async def _generate_cluster_summary(
    facts_text_list: list[str],
    cluster_label: str,
    llm: LLMProvider,
) -> str:
    """Call LLM to summarize facts in 2-3 sentences."""
    if not facts_text_list:
        return ""
    system_prompt = (
        "You are a personal memory system. Summarize the facts below in 2-3 short sentences.\n\n"
        "CRITICAL RULES:\n"
        "1. Be factual and concise. Do NOT invent information not in the facts.\n"
        "2. Do NOT make calculations or mathematical inferences.\n"
        "3. Each fact has format '[entity:key] attribute = value'. Respect attribution.\n"
        "4. For facts about 'user:self', refer to them as 'The user'.\n"
        "5. Do NOT interpret numbers out of context.\n"
        "6. If facts are about different entities, clearly separate who is who.\n"
        "7. Write the summary in the same language as the fact content."
    )
    facts_block = "\n".join(f"- {t}" for t in facts_text_list[:50])
    user_content = f"Cluster: {cluster_label}\n\nFacts:\n{facts_block}"
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": user_content},
    ]
    try:
        _llm_result = await llm.complete(messages)
        if _llm_result.text:
            return _llm_result.text.strip()[:2000]
    except Exception as e:
        logger.debug("cluster_summary_generation_failed: %s", e)
    return ""


# ---------------------------------------------------------------------------
# Main clustering
# ---------------------------------------------------------------------------


async def cluster_user_facts(
    session: AsyncSession,
    user_id: str,
    embedding_provider: EmbeddingProvider,
    llm_provider: LLMProvider,
    config: MemoryConfig,
) -> ClusteringResult:
    """Build or update clusters for a user's facts. Idempotent.

    Groups active facts by ``(entity_type, entity_key)``, creates or updates
    clusters with LLM-generated summaries and embeddings.

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

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

    # Load active facts within age window
    stmt = (
        select(MemoryFact)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.valid_to.is_(None),
            MemoryFact.valid_from >= cutoff,
        )
        .order_by(MemoryFact.valid_from.desc())
        .options(defer(MemoryFact.embedding))
    )
    result = await session.execute(stmt)
    facts = list(result.scalars().all())

    # Group by (entity_type, entity_key)
    groups: dict[tuple[str, str], list[Any]] = defaultdict(list)
    for f in facts:
        groups[(f.entity_type or "", f.entity_key or "")].append(f)

    # Load existing active clusters
    existing_stmt = select(MemoryCluster).where(MemoryCluster.user_id == user_id, MemoryCluster.is_active.is_(True))
    existing_result = await session.execute(existing_stmt)
    existing_clusters = list(existing_result.scalars().all())

    # Map (entity_type, entity_key) → cluster via representative facts
    cluster_by_entity: dict[tuple[str, str], Any] = {}
    if existing_clusters:
        cluster_ids_list = [c.id for c in existing_clusters]
        repr_stmt = (
            select(
                MemoryFact.cluster_id,
                MemoryFact.entity_type,
                MemoryFact.entity_key,
            )
            .where(MemoryFact.cluster_id.in_(cluster_ids_list))
            .distinct(MemoryFact.cluster_id)
            .order_by(MemoryFact.cluster_id, MemoryFact.valid_from.desc())
        )
        repr_result = await session.execute(repr_stmt)
        cluster_map_by_id = {c.id: c for c in existing_clusters}
        for row in repr_result.all():
            cid, etype, ekey = row[0], row[1], row[2]
            if cid in cluster_map_by_id:
                cluster_by_entity[(etype or "", ekey or "")] = cluster_map_by_id[cid]

    stats = ClusteringResult()

    # Unassign facts from clusters for reassignment
    fact_ids = [f.id for f in facts]
    if fact_ids:
        await session.execute(update(MemoryFact).where(MemoryFact.id.in_(fact_ids)).values(cluster_id=None))

    for (entity_type, entity_key), group_facts in groups.items():
        if len(group_facts) < config.cluster_min_facts:
            continue

        label = _cluster_label_for_entity(entity_type, entity_key)
        facts_text_list: list[str] = []
        for f in group_facts:
            attr = f.attribute_key or "unknown"
            val = f.value_text or str(f.value_json)
            entity = f"{f.entity_type}:{f.entity_key}" if f.entity_type else "unknown"
            if val:
                facts_text_list.append(f"[{entity}] {attr} = {val}"[:500])

        summary_text = await _generate_cluster_summary(facts_text_list, label, llm_provider)
        importance_avg = sum(f.importance or 0.5 for f in group_facts) / len(group_facts)

        cluster = cluster_by_entity.get((entity_type, entity_key))
        if cluster is None:
            cluster = MemoryCluster(
                id=uuid.uuid4(),
                user_id=user_id,
                label=label,
                summary_text=summary_text or None,
                cluster_type="auto",
                fact_count=len(group_facts),
                importance=round(importance_avg, 4),
                is_active=True,
            )
            session.add(cluster)
            await session.flush()
            stats.clusters_created += 1
        else:
            cluster.label = label
            cluster.summary_text = summary_text or cluster.summary_text
            cluster.fact_count = len(group_facts)
            cluster.importance = round(importance_avg, 4)
            cluster.last_updated_at = datetime.now(UTC)
            stats.clusters_reinforced += 1

        if summary_text:
            stats.summaries_generated += 1
            emb = await embedding_provider.embed_one(summary_text)
            if emb:
                cluster.embedding_vec = emb

        for f in group_facts:
            f.cluster_id = cluster.id
            stats.facts_assigned += 1

    # Deactivate clusters that have no facts
    if existing_clusters:
        cluster_ids_for_count = [c.id for c in existing_clusters]
        count_stmt = (
            select(MemoryFact.cluster_id, func.count().label("cnt"))
            .where(MemoryFact.cluster_id.in_(cluster_ids_for_count))
            .group_by(MemoryFact.cluster_id)
        )
        count_result = await session.execute(count_stmt)
        count_map: dict[Any, int] = {row[0]: row[1] for row in count_result.all()}
        for c in existing_clusters:
            if count_map.get(c.id, 0) == 0:
                c.is_active = False

    logger.info(
        "cluster_user_facts: user=%s, created=%d, reinforced=%d, summaries=%d, facts=%d",
        user_id,
        stats.clusters_created,
        stats.clusters_reinforced,
        stats.summaries_generated,
        stats.facts_assigned,
    )
    return stats


# ---------------------------------------------------------------------------
# Incremental cluster rebuild for corrected facts
# ---------------------------------------------------------------------------


async def rebuild_clusters_for_corrected_facts(
    session: AsyncSession,
    user_id: str,
    corrected_fact_ids: list[str],
    llm_provider: LLMProvider,
    embedding_provider: EmbeddingProvider,
) -> int:
    """Rebuild only clusters affected by corrected facts.

    More efficient than full rebuild — identifies clusters containing
    the corrected facts and regenerates their summaries.

    Args:
        session: Database session.
        user_id: User identifier.
        corrected_fact_ids: IDs of facts that were corrected.
        llm_provider: Injected LLM provider.
        embedding_provider: Injected embedding provider.

    Returns:
        Number of clusters rebuilt.
    """
    if not corrected_fact_ids:
        return 0

    # Find affected clusters
    affected_stmt = (
        select(MemoryFact.cluster_id)
        .where(
            MemoryFact.user_id == user_id,
            MemoryFact.id.in_(corrected_fact_ids),
            MemoryFact.cluster_id.isnot(None),
        )
        .distinct()
    )
    result = await session.execute(affected_stmt)
    cluster_ids = [row[0] for row in result.all() if row[0] is not None]

    if not cluster_ids:
        return 0

    # Load affected clusters
    clusters_stmt = select(MemoryCluster).where(MemoryCluster.id.in_(cluster_ids), MemoryCluster.is_active.is_(True))
    clusters_result = await session.execute(clusters_stmt)
    clusters = list(clusters_result.scalars().all())

    rebuilt = 0
    for cluster in clusters:
        # Fetch current facts for this cluster
        facts_stmt = (
            select(MemoryFact)
            .where(
                MemoryFact.cluster_id == cluster.id,
                MemoryFact.valid_to.is_(None),
            )
            .order_by(MemoryFact.valid_from.desc())
            .options(defer(MemoryFact.embedding))
        )
        facts_result = await session.execute(facts_stmt)
        facts = list(facts_result.scalars().all())

        if not facts:
            cluster.is_active = False
            continue

        # Regenerate summary
        facts_text_list = [
            f"[{f.entity_type}:{f.entity_key}] {f.attribute_key} = {f.value_text or str(f.value_json)}"[:500]
            for f in facts
            if f.value_text or f.value_json
        ]
        summary = await _generate_cluster_summary(facts_text_list, cluster.label or "", llm_provider)
        if summary:
            cluster.summary_text = summary
            cluster.fact_count = len(facts)
            cluster.last_updated_at = datetime.now(UTC)
            emb = await embedding_provider.embed_one(summary)
            if emb:
                cluster.embedding_vec = emb
            rebuilt += 1

    logger.info(
        "rebuild_clusters_for_corrected_facts: user=%s, affected=%d, rebuilt=%d",
        user_id,
        len(clusters),
        rebuilt,
    )
    return rebuilt


# ---------------------------------------------------------------------------
# Union-Find for community detection
# ---------------------------------------------------------------------------


class _UnionFind:
    """Simple Union-Find for grouping cluster 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


# ---------------------------------------------------------------------------
# Community detection
# ---------------------------------------------------------------------------

DUPLICATE_SIMILARITY_THRESHOLD = 0.9


async def detect_communities(
    session: AsyncSession,
    user_id: str,
    embedding_provider: EmbeddingProvider,
    llm_provider: LLMProvider,
    config: MemoryConfig,
) -> CommunityDetectionResult:
    """Find cross-entity themes by comparing cluster embeddings.

    Uses Union-Find to group clusters from different entity_keys with
    similar embeddings into communities, materialized as MemoryMetaObservation.

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

    Returns:
        CommunityDetectionResult with stats.
    """
    threshold = config.community_similarity_threshold
    min_clusters = config.community_min_clusters

    # Load active clusters with embeddings
    stmt = select(MemoryCluster).where(
        MemoryCluster.user_id == user_id,
        MemoryCluster.is_active.is_(True),
        MemoryCluster.embedding_vec.isnot(None),
    )
    result = await session.execute(stmt)
    clusters = list(result.scalars().all())

    if len(clusters) < min_clusters:
        return CommunityDetectionResult(skipped=True, skip_reason="insufficient_clusters")

    # Map cluster → entity_key via representative fact
    cluster_ids_list = [c.id for c in clusters]
    repr_stmt = (
        select(MemoryFact.cluster_id, MemoryFact.entity_type, MemoryFact.entity_key)
        .where(MemoryFact.cluster_id.in_(cluster_ids_list))
        .distinct(MemoryFact.cluster_id)
        .order_by(MemoryFact.cluster_id, MemoryFact.valid_from.desc())
    )
    repr_result = await session.execute(repr_stmt)
    entity_map: dict[Any, tuple[str, str]] = {}
    for row in repr_result.all():
        entity_map[row[0]] = (row[1] or "", row[2] or "")

    indexed_clusters = [(i, c) for i, c in enumerate(clusters) if c.id in entity_map]
    if len(indexed_clusters) < min_clusters:
        return CommunityDetectionResult(skipped=True, skip_reason="insufficient_mapped_clusters")

    # Pairwise similarity + Union-Find
    n = len(indexed_clusters)
    uf = _UnionFind(n)
    for i in range(n):
        ci = indexed_clusters[i][1]
        ei = entity_map[ci.id]
        vi = list(ci.embedding_vec)
        for j in range(i + 1, n):
            cj = indexed_clusters[j][1]
            ej = entity_map[cj.id]
            if ei == ej:
                continue
            vj = list(cj.embedding_vec)
            sim = cosine_similarity(vi, vj)
            if sim >= threshold:
                uf.union(i, j)

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

    if not communities:
        return CommunityDetectionResult()

    # Load existing community_theme observations for dedup
    existing_stmt = select(MemoryMetaObservation).where(
        MemoryMetaObservation.user_id == user_id,
        MemoryMetaObservation.observation_type == "community_theme",
        MemoryMetaObservation.is_active.is_(True),
    )
    existing_result = await session.execute(existing_stmt)
    existing_obs = list(existing_result.scalars().all())

    result_obj = CommunityDetectionResult()

    for members in communities:
        member_clusters = [indexed_clusters[m][1] for m in members]
        cluster_summaries = [c.summary_text for c in member_clusters if c.summary_text]
        if not cluster_summaries:
            continue

        # LLM title + summary
        title = await _generate_community_title(cluster_summaries, llm_provider)
        if not title:
            continue
        text = await _generate_community_summary(cluster_summaries, llm_provider)

        # Embed title for dedup
        title_embedding: list[float] | None = None
        try:
            title_embedding = await embedding_provider.embed_one(title)
        except Exception as e:
            logger.debug("community_title_embedding_failed: %s", e)

        # Dedup
        duplicate = None
        if title_embedding and existing_obs:
            for ex in existing_obs:
                if ex.embedding_vec is None:
                    continue
                if cosine_similarity(title_embedding, list(ex.embedding_vec)) >= DUPLICATE_SIMILARITY_THRESHOLD:
                    duplicate = ex
                    break

        if duplicate:
            duplicate.times_reinforced = (duplicate.times_reinforced or 1) + 1
            duplicate.last_reinforced_at = datetime.now(UTC)
            result_obj.communities_reinforced += 1
        else:
            new_obs = MemoryMetaObservation(
                user_id=user_id,
                observation_type="community_theme",
                title=title,
                text=text or "",
                supporting_event_ids=[],
                supporting_fact_ids=[],
                confidence=0.7,
                importance=0.8,
                times_reinforced=1,
                is_active=True,
                embedding_vec=title_embedding,
            )
            session.add(new_obs)
            existing_obs.append(new_obs)
            result_obj.communities_created += 1

        result_obj.clusters_in_communities += len(members)

    logger.info(
        "detect_communities: user=%s, created=%d, reinforced=%d, clusters=%d",
        user_id,
        result_obj.communities_created,
        result_obj.communities_reinforced,
        result_obj.clusters_in_communities,
    )
    return result_obj


async def _generate_community_summary(cluster_summaries: list[str], llm: LLMProvider) -> str:
    """LLM-generated summary for a cross-entity community."""
    if not cluster_summaries:
        return ""
    system_prompt = (
        "You are a personal memory system. The texts below are summaries of "
        "thematic clusters from DIFFERENT entities of the same user. "
        "Identify the CROSS-CUTTING THEME that connects these clusters and summarize in "
        "2-3 short sentences. Be factual and concise. Write in the same language as the content."
    )
    block = "\n---\n".join(cluster_summaries[:10])
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": block},
    ]
    try:
        _llm_result = await llm.complete(messages)
        if _llm_result.text:
            return _llm_result.text.strip()[:2000]
    except Exception as e:
        logger.debug("community_summary_generation_failed: %s", e)
    return ""


async def _generate_community_title(cluster_summaries: list[str], llm: LLMProvider) -> str:
    """LLM-generated title for a cross-entity community (max 80 chars)."""
    if not cluster_summaries:
        return ""
    system_prompt = (
        "Generate a short title (max 80 characters) that captures the cross-cutting theme "
        "between the clusters below. Write in the same language as the content. "
        "Return only the title, no quotes or explanation."
    )
    block = "\n---\n".join(cluster_summaries[:10])
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": block},
    ]
    try:
        _llm_result = await llm.complete(messages)
        if _llm_result.text:
            return _llm_result.text.strip()[:80]
    except Exception as e:
        logger.debug("community_title_generation_failed: %s", e)
    return ""
