"""v0.2 enhanced consolidation engine with LLM integration."""

from __future__ import annotations

import logging
import time
from collections import defaultdict

from neurovault.cognitive.consolidation import ConsolidationEngine
from neurovault.cognitive.importance_calibrator import ImportanceCalibrator
from neurovault.cognitive.llm_client import OllamaClient
from neurovault.cognitive.llm_insights import LLMInsightGenerator
from neurovault.cognitive.llm_relationships import LLMRelationshipExtractor
from neurovault.models import ConsolidationReport

logger = logging.getLogger(__name__)


class ConsolidationEngineV2(ConsolidationEngine):
    """v0.2 consolidation with LLM-enhanced phases."""

    def __init__(self, storage, chronos_bridge, extractor=None, llm_client=None, calibrator=None):
        super().__init__(storage, chronos_bridge, extractor)
        self._llm = llm_client or OllamaClient()
        self._insight_gen = LLMInsightGenerator(self._llm)
        self._rel_extractor = LLMRelationshipExtractor(self._llm)
        self._calibrator = calibrator or ImportanceCalibrator()
        self._llm_available = self._llm.is_available()

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

        if "entities" in all_phases:
            report.entities_extracted, report.relationships_found = self._phase1_entities()

        if "llm_relations" in all_phases and self._llm_available:
            report.relationships_found += self._phase1b_llm_relationships()

        if "clusters" in all_phases:
            report.clusters_formed = self._phase2_clusters()

        if "insights" in all_phases:
            report.insights_generated = self._phase3_insights_v2()

        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()

        if "calibrate" in all_phases:
            self._phase6_calibrate()

        report.duration_seconds = time.time() - start
        self._bridge.checkpoint(
            f"v0.2 Consolidation: {report.entities_extracted} entities, "
            f"{report.relationships_found} rels, {report.insights_generated} insights, "
            f"{report.memories_decayed} decayed [LLM: {'on' if self._llm_available else 'off'}]"
        )
        return report

    def _phase1b_llm_relationships(self) -> int:
        """Extract typed relationships using LLM."""
        count = 0
        unprocessed = self._storage.get_unprocessed_memory_ids("llm_relationships", limit=30)

        for mid in unprocessed:
            try:
                mem = self._bridge.get_memory(mid)
            except Exception:
                continue

            triples = self._rel_extractor.extract(mem.content)
            for triple in triples:
                self._storage.upsert_relationship(
                    source_name=triple["subject"],
                    target_name=triple["object"],
                    relation_type=triple["predicate"],
                    strength=triple["confidence"],
                )
                count += 1

            self._storage.mark_processed(mid, "llm_relationships")

        return count

    def _phase3_insights_v2(self) -> int:
        """LLM-powered insight generation (falls back to template)."""
        clusters = self._storage.get_clusters_without_insights()
        insights = 0

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

            texts, entity_names = self._gather_cluster_data(cluster)
            insight_text = self._insight_gen.generate(
                memory_texts=texts,
                entity_names=entity_names,
                cluster_size=len(cluster.memory_ids),
            )

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

        return insights

    def _phase6_calibrate(self) -> None:
        """Auto-calibrate importance scores based on access patterns."""

    def _gather_cluster_data(self, cluster) -> tuple[list[str], list[str]]:
        entity_counts: dict[str, int] = defaultdict(int)
        texts: list[str] = []
        for mid in cluster.memory_ids[:10]:
            entities = self._storage.get_entities_for_memory(mid)
            for e in entities:
                entity_counts[e.name] += 1
            try:
                mem = self._bridge.get_memory(mid)
                texts.append(mem.content[:200])
            except Exception as exc:
                logger.debug("Cluster data gather skipped memory %s: %s", mid.hex(), exc)
        top = sorted(entity_counts.items(), key=lambda x: -x[1])[:5]
        return texts, [e[0] for e in top]
