"""EpisodeDigestor — reads conversation_log, distills knowledge, ingests memories.

Two modes (configured via DIGESTOR_MODE env var):
- "distill" (default): Gemini distills turns into clean knowledge entries
- "segment" (legacy): Groq segments by topic, stores raw episodes

Flow (distill mode):
- run(): main loop — processes ALL undigested data, sleeps only when idle
- Within each cycle: fetch all entries → prepare turns → split into ~32k token batches
- Each batch → Gemini distills → ingest each knowledge entry via cortex
- When queue is empty → sleep 30s → recheck

Flow (segment mode — legacy):
- Same batching, but Groq segments by topic → raw episodes stored
"""

import logging
import re
from typing import TYPE_CHECKING

from neo_cortex import config
from neo_cortex.models import ConversationEntry
from neo_cortex.reflection import meta_reflect

if TYPE_CHECKING:
    from neo_cortex.classifier import Classifier
    from neo_cortex.conversation_log import ConversationLog
    from neo_cortex.cortex import CortexService
    from neo_cortex.distiller import GeminiDistiller
    from neo_cortex.merge_consolidator import MemoryConsolidator

logger = logging.getLogger("neo-cortex-digestor")

MAX_EPISODE_CHARS = 8000
CHARS_PER_TOKEN = 4

# Re-export for backward compat with tests
TOKEN_BUDGET = config.DIGESTOR_BUDGET


from functools import partial
from neo_cortex._log import boot_log
_bl = partial(boot_log, "DIGESTOR")


class EpisodeDigestor:
    """Background worker: processes all undigested conversation data continuously.

    Supports two modes:
    - distill: GeminiDistiller → clean knowledge entries (preferred)
    - segment: Classifier → raw topic segments (legacy fallback)
    """

    def __init__(
        self,
        log: 'ConversationLog | None',
        cortex: 'CortexService',
        classifier: 'Classifier | None' = None,
        distiller: 'GeminiDistiller | None' = None,
        consolidator: 'MemoryConsolidator | None' = None,
    ):
        self._log = log
        self._cortex = cortex
        self._classifier = classifier
        self._distiller = distiller
        self._consolidator = consolidator
        self._total_ingested = 0
        self._paused = False
        self._needs_dream = False

    def pause(self) -> None:
        """Pause processing (for memory_rebuild)."""
        self._paused = True
        _bl("digestor PAUSED")

    def resume(self) -> None:
        """Resume processing after rebuild."""
        self._paused = False
        _bl("digestor RESUMED")

    async def run(self) -> None:
        """Main loop — process everything, sleep only when idle."""
        import asyncio

        mode = "distill" if self._distiller else "segment"
        _bl(f"run: started (mode={mode})")
        while True:
            if self._paused:
                await asyncio.sleep(5)
                continue
            try:
                did_work = await self._process_all()
            except Exception as e:
                _bl(f"run: cycle FAILED: {e}")
                did_work = False

            if did_work:
                self._needs_dream = True

            if not did_work:
                # Idle — run consolidation if available
                if self._consolidator:
                    try:
                        merged = await self._consolidator.consolidate()
                        if merged:
                            _bl(f"run: consolidation merged {merged} groups")
                    except Exception as e:
                        _bl(f"run: consolidation FAILED: {e}")
                # Dream after consolidation (if work was done before)
                if self._needs_dream:
                    try:
                        dream_result = await self._cortex.dream()
                        _bl(f"run: dream completed — surprisal={len(dream_result.surprisal)}")
                    except Exception as e:
                        _bl(f"run: dream FAILED: {e}")
                    self._needs_dream = False
                    # Meta-cognitive reflection after dream
                    try:
                        from neo_cortex.llm import create_client
                        llm = create_client()
                        n_insights = await meta_reflect(self._cortex, llm)
                        _bl(f"run: meta_reflect completed — {n_insights} insights")
                    except Exception as e:
                        _bl(f"run: meta_reflect FAILED (non-blocking): {e}")
                await asyncio.sleep(30)

    async def _process_all(self) -> bool:
        """Fetch all undigested, process in token-budgeted batches. Returns True if work done."""
        if not self._log:
            return False

        entries = self._log.get_undigested(limit=10000)
        if not entries:
            return False

        _bl(f"processing: {len(entries)} undigested entries")

        turns = self._prepare_turns(entries)
        if not turns:
            self._log.mark_digested([e.id for e in entries])
            _bl("processing: all entries were empty/noise, marked digested")
            return True

        _bl(f"processing: {len(turns)} turns prepared")

        if self._distiller:
            await self._process_distill(turns, entries)
        else:
            await self._process_segment(turns)

        _bl(f"processing: done — total ingested this cycle: {self._total_ingested}")
        return True

    # ── Distill mode (Gemini) ──────────────────────────────────────────

    async def _process_distill(self, turns: list[dict], entries: list[ConversationEntry]) -> None:
        """Process turns via Gemini distillation — no carryover needed."""
        all_entry_ids = [e.id for e in entries]
        pos = 0

        # Dynamic budget: large backlog → 100k tokens, regular flow → 32k
        budget = config.DIGESTOR_BUDGET_MAX if len(turns) > config.DIGESTOR_BUDGET_THRESHOLD else config.DIGESTOR_BUDGET
        _bl(f"distill: {len(turns)} turns, budget={budget} tokens ({'max' if budget == config.DIGESTOR_BUDGET_MAX else 'normal'})")

        while pos < len(turns):
            # Build batch up to token budget
            batch_texts: list[str] = []
            batch_entry_ids: list[int] = []
            char_budget = budget * CHARS_PER_TOKEN
            session_id = turns[pos].get("session_id", "unknown")

            while pos < len(turns) and char_budget > 0:
                t = turns[pos]
                batch_texts.append(t["full_text"])
                batch_entry_ids.extend(t["entry_ids"])
                char_budget -= len(t["full_text"])
                pos += 1

            if not batch_texts:
                break

            _bl(f"distill batch: {len(batch_texts)} turns, ~{sum(len(t) for t in batch_texts) // 4} tokens")

            # Gemini distill
            try:
                knowledge_entries = await self._distiller.distill(batch_texts)
            except Exception as e:
                _bl(f"distill: FAILED: {e}, marking batch as digested anyway")
                knowledge_entries = []

            # Ingest each knowledge entry (fast path — no conflict resolution)
            for ke in knowledge_entries:
                try:
                    mid = await self._cortex.ingest_knowledge_fast(
                        topic=ke.topic,
                        content=ke.content,
                        entry_type=ke.type,
                        project=ke.project,
                        session_id=session_id,
                        concepts=ke.concepts if ke.concepts else None,
                    )
                    if mid:
                        self._total_ingested += 1
                        _bl(f"ingested: '{ke.topic}' [{ke.project}] → {mid}")
                    else:
                        _bl(f"filtered: '{ke.topic}' (dedup/conflict)")
                except Exception as e:
                    _bl(f"FAILED ingest: '{ke.topic}' error={e}")

            # Mark batch entries as digested
            self._log.mark_digested(batch_entry_ids)

    # ── Segment mode (Groq legacy) ────────────────────────────────────

    async def _process_segment(self, turns: list[dict]) -> None:
        """Process turns via Groq topic segmentation with carryover."""
        carryover: list[dict] = []
        pos = 0

        while pos < len(turns) or carryover:
            batch = list(carryover)
            carryover = []
            char_budget = config.DIGESTOR_BUDGET * CHARS_PER_TOKEN

            for t in batch:
                char_budget -= len(t["text"])

            pos_before = pos
            while pos < len(turns) and char_budget > 0:
                t = turns[pos]
                batch.append(t)
                char_budget -= len(t["text"])
                pos += 1

            if not batch:
                break

            added_new = (pos > pos_before)
            is_last_batch = (pos >= len(turns)) or not added_new

            segments = await self._segment(batch)

            if not segments:
                all_ids = []
                for t in batch:
                    all_ids.extend(t["entry_ids"])
                self._log.mark_digested(all_ids)
                continue

            if is_last_batch:
                for seg in segments:
                    await self._ingest_segment(seg, batch)
            else:
                for seg in segments[:-1]:
                    await self._ingest_segment(seg, batch)

                last = segments[-1]
                start = max(0, min(last.get("start", 0), len(batch) - 1))
                end = max(start, min(last.get("end", len(batch) - 1), len(batch) - 1))
                carryover = batch[start:end + 1]
                _bl(f"carryover: {len(carryover)} turns from topic '{last.get('topic', '?')}'")

    async def _segment(self, batch: list[dict]) -> list[dict]:
        """Segment a batch of turns by topic via Groq."""
        if not self._classifier or len(batch) <= 1:
            return [{"start": 0, "end": len(batch) - 1, "topic": "unknown"}]

        try:
            segments = await self._classifier.segment_topics(
                [t["text"] for t in batch]
            )
        except Exception as e:
            _bl(f"segment: Groq FAILED: {e}, fallback to chunks of 5")
            segments = [
                {"start": i, "end": min(i + 4, len(batch) - 1), "topic": "unknown"}
                for i in range(0, len(batch), 5)
            ]

        return segments or []

    async def _ingest_segment(self, seg: dict, batch: list[dict]) -> None:
        """Ingest one topic segment as an episode."""
        start = seg.get("start", 0)
        end = seg.get("end", len(batch) - 1)
        topic = seg.get("topic", "unknown")

        start = max(0, min(start, len(batch) - 1))
        end = max(start, min(end, len(batch) - 1))

        segment_entry_ids = []
        text_parts = []
        for t in batch[start:end + 1]:
            segment_entry_ids.extend(t["entry_ids"])
            text_parts.append(t["full_text"])

        text = "\n".join(text_parts)
        if len(text) > MAX_EPISODE_CHARS:
            text = text[:MAX_EPISODE_CHARS]

        if not text.strip():
            self._log.mark_digested(segment_entry_ids)
            return

        session_id = batch[start].get("session_id", "unknown")

        try:
            mid = await self._cortex.ingest_episode(text, session_id)
            if mid:
                self._total_ingested += 1
                _bl(f"ingested: topic='{topic}' turns={start}-{end} chars={len(text)} → {mid}")
            else:
                _bl(f"filtered: topic='{topic}' turns={start}-{end} (dedup/empty)")
        except Exception as e:
            _bl(f"FAILED: topic='{topic}' turns={start}-{end} error={e}")

        self._log.mark_digested(segment_entry_ids)

    # ── Shared ─────────────────────────────────────────────────────────

    def _prepare_turns(self, entries: list[ConversationEntry]) -> list[dict]:
        """Clean entries and return as individual turns for Gemini distillation.

        Same approach as the manual test that produced best results:
        - Filter to user/assistant only
        - Strip thinking tags and XML
        - Skip entries < 20 chars after cleanup
        - Each entry is a separate turn
        """
        _xml_re = re.compile(r"<[^>]+>")

        turns: list[dict] = []
        for e in entries:
            if e.role not in ("user", "assistant"):
                continue

            text = _xml_re.sub("", e.content).strip()
            if not text:
                continue

            label = "USER" if e.role == "user" else "ASSISTANT"
            full_text = f"{label}: {text}"

            turns.append({
                "text": full_text[:500],
                "full_text": full_text,
                "entry_ids": [e.id],
                "session_id": e.session_id,
            })

        return turns
