"""5-phase brain-inspired consolidation engine."""

from __future__ import annotations

import math
import time
from collections import defaultdict

try:
    import numpy as np
except Exception:  # pragma: no cover
    np = None  # type: ignore

from neurovault.cognitive.entity_extractor import EntityExtractor
from neurovault.models import ConsolidationReport
from neurovault.storage import NeurovaultStorage


class ConsolidationEngine:
    """Runs during idle periods - NEUROVAULT's sleep cycle."""

    def __init__(self, storage: NeurovaultStorage, chronos_bridge, extractor: EntityExtractor | None = None):
        self._storage = storage
        self._bridge = chronos_bridge
        self._extractor = extractor or EntityExtractor()

    def run(self, force: bool = False, phases: list[str] | None = None) -> ConsolidationReport:
        del force
        start = time.time()
        report = ConsolidationReport()
        all_phases = phases or ["entities", "clusters", "insights", "promote", "decay"]

        if "entities" in all_phases:
            report.entities_extracted, report.relationships_found = self._phase1_entities()
        if "clusters" in all_phases:
            report.clusters_formed = self._phase2_clusters()
        if "insights" in all_phases:
            report.insights_generated = self._phase3_insights()
        if "promote" in all_phases:
            report.memories_promoted = self._phase4_promotion()
        if "decay" in all_phases:
            report.memories_decayed, report.memories_archived = self._phase5_decay()

        report.duration_seconds = time.time() - start
        self._bridge.checkpoint(
            f"Consolidation: {report.entities_extracted} entities, "
            f"{report.insights_generated} insights, {report.memories_decayed} decayed"
        )
        return report

    def _phase1_entities(self) -> tuple[int, int]:
        entity_count = 0
        relationship_count = 0
        unprocessed_ids = self._storage.get_unprocessed_memory_ids("entity_extraction", limit=300)

        for memory_id in unprocessed_ids:
            try:
                memory = self._bridge.get_memory(memory_id)
            except Exception:
                continue

            raw_entities = self._extractor.extract(memory.content)
            for raw in raw_entities:
                entity_type = raw["type"].value if hasattr(raw["type"], "value") else raw["type"]
                entity = self._storage.upsert_entity(
                    name=raw["name"],
                    entity_type=entity_type,
                    confidence=raw.get("confidence", 0.8),
                    seen_at=memory.created_at,
                )
                self._storage.link_memory_entity(memory.id, entity.id, raw.get("confidence", 0.8))
                entity_count += 1

            if len(raw_entities) >= 2:
                for i, left in enumerate(raw_entities):
                    for right in raw_entities[i + 1 :]:
                        self._storage.upsert_relationship(
                            source_name=left["name"],
                            target_name=right["name"],
                            relation_type=self._infer_relation(left, right),
                            strength=0.3,
                        )
                        relationship_count += 1

            self._storage.mark_processed(memory.id, "entity_extraction")

        return entity_count, relationship_count

    def _phase2_clusters(self) -> int:
        if np is None:
            return 0
        memories = self._bridge.get_recent_memories_with_embeddings(limit=300)
        if len(memories) < 5:
            return 0

        embeddings = np.array([m["embedding"] for m in memories], dtype=np.float32)
        if embeddings.shape[0] < 5:
            return 0

        norms = np.linalg.norm(embeddings, axis=1, keepdims=True)
        normalized = embeddings / (norms + 1e-8)
        similarity_matrix = normalized @ normalized.T

        threshold = 0.6
        assigned = set()
        clusters_formed = 0

        for i in range(len(memories)):
            if i in assigned:
                continue
            cluster = [i]
            assigned.add(i)
            for j in range(i + 1, len(memories)):
                if j in assigned:
                    continue
                if similarity_matrix[i][j] >= threshold:
                    cluster.append(j)
                    assigned.add(j)

            if len(cluster) >= 3:
                memory_ids = [memories[idx]["id"] for idx in cluster]
                centroid = np.mean([embeddings[idx] for idx in cluster], axis=0)
                self._storage.upsert_cluster(memory_ids, centroid)
                clusters_formed += 1

        return clusters_formed

    def _phase3_insights(self) -> int:
        clusters = self._storage.get_clusters_without_insights()
        insights_generated = 0

        for cluster in clusters:
            if len(cluster.memory_ids) < 3:
                continue

            entity_counts: dict[str, int] = defaultdict(int)
            memory_summaries: list[str] = []

            for memory_id in cluster.memory_ids[:10]:
                entities = self._storage.get_entities_for_memory(memory_id)
                for entity in entities:
                    entity_counts[entity.name] += 1
                try:
                    memory = self._bridge.get_memory(memory_id)
                    memory_summaries.append(memory.content[:80])
                except Exception:
                    continue

            top_entities = sorted(entity_counts.items(), key=lambda kv: -kv[1])[:5]
            entity_names = [name for name, _ in top_entities]
            insight_text = self._generate_insight(entity_names, memory_summaries[:3], len(cluster.memory_ids))

            insight = self._bridge.store_memory(
                content=insight_text,
                source="consolidation",
                importance=0.8,
                metadata={
                    "type": "insight",
                    "cluster_id": cluster.id.hex() if cluster.id else "",
                    "source_count": len(cluster.memory_ids),
                    "key_entities": entity_names,
                },
            )
            self._storage.set_cluster_insight(cluster.id, insight.id)
            insights_generated += 1

        return insights_generated

    def _phase4_promotion(self) -> int:
        promoted = 0
        promoted += self._storage.promote_tier("working", "short_term", min_age_seconds=30)
        promoted += self._storage.promote_high_value("short_term", "long_term")
        return promoted

    def _phase5_decay(self) -> tuple[int, int]:
        decayed = 0
        archived = 0

        memories = self._bridge.get_active_memories_with_metadata()
        for memory in memories:
            metadata = memory.get("metadata", {}) or {}
            tier = metadata.get("tier", "long_term")
            if tier == "archive":
                continue

            access_count = int(metadata.get("access_count", 0))
            last_accessed = int(metadata.get("last_accessed", memory.get("created_at", 0)))
            hours_since = (time.time() * 1000 - last_accessed) / (3600 * 1000)

            original_importance = float(memory.get("importance", 0.5))
            decayed_importance = self._ebbinghaus(original_importance, hours_since, access_count)

            if decayed_importance < original_importance * 0.5:
                self._bridge.update_importance(memory["id"], decayed_importance)
                decayed += 1

                if decayed_importance < 0.05 and tier != "archive":
                    self._bridge.update_memory_metadata(memory["id"], {"tier": "archive"})
                    archived += 1

        return decayed, archived

    @staticmethod
    def _ebbinghaus(importance: float, hours_since_last_access: float, access_count: int) -> float:
        stability = 1.0 * (1 + importance * 2)
        reinforcement = max(1, access_count)
        effective_stability = stability * math.log2(reinforcement + 1)
        if effective_stability <= 0:
            return 0.01
        retention = math.exp(-hours_since_last_access / (effective_stability * 24))
        return max(importance * retention, 0.01)

    @staticmethod
    def _infer_relation(left: dict, right: dict) -> str:
        t1 = left["type"].value if hasattr(left["type"], "value") else left["type"]
        t2 = right["type"].value if hasattr(right["type"], "value") else right["type"]
        pair = frozenset([t1, t2])
        mapping = {
            frozenset(["PERSON", "PROJECT"]): "works-on",
            frozenset(["PROJECT", "TECHNOLOGY"]): "uses",
            frozenset(["PROJECT", "DATE"]): "has-deadline",
            frozenset(["PERSON", "ORGANIZATION"]): "works-at",
            frozenset(["PERSON", "PERSON"]): "collaborates-with",
        }
        return mapping.get(pair, "related-to")

    @staticmethod
    def _generate_insight(entities: list[str], summaries: list[str], count: int) -> str:
        entity_str = ", ".join(entities[:3]) if entities else "various topics"
        summary_str = "; ".join(summaries[:2]) if summaries else ""
        return f"Knowledge pattern ({count} connected memories): {entity_str}. Key observations: {summary_str}".strip()
