"""
MODULE 6c: CURIOSITY SEEDS + CROSS-MEMORY CONNECTIONS
======================================================
Implements the generative intelligence layer of the dream cycle:
  - Curiosity seed extraction (2-5 open questions / patterns per cycle)
  - Thematic clustering (group memories by entity/topic overlap)
  - Cross-memory edge creation (RELATED_TO edges in the memory graph)
  - Dream journal generation (human-readable markdown summary)

All public methods are async. Hard dependencies injected (no global state).

This module is called from SleepConsolidator.dream() at step 9 (DREAM JOURNAL)
after all consolidation, decay, and graph refresh steps have run.

Integration:
  - CuriosityEngine is instantiated by SleepConsolidator and called at the end
    of the dream() pipeline.
  - CuriosityReport is embedded into DreamReport.journal and optionally
    stored in session_record.dream_report as JSON.
  - Seeds with memory_type='curiosity_seed' are stored as SemanticClaims so
    they surface at the next wake() call.
"""

from __future__ import annotations

import json
import logging
import time
import uuid
from dataclasses import dataclass, field
from datetime import datetime, timezone
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple

if TYPE_CHECKING:
    from cortex.storage.adapter import StorageAdapter
    from cortex.config import CortexConfig
    from cortex.types import LLMProvider, SemanticClaim

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# LLM prompt templates
# ---------------------------------------------------------------------------

_CURIOSITY_SYSTEM = """\
You are a reflective memory analyst for an AI assistant. You have just reviewed
the memories from a completed session. Your job is to identify 2-5 "curiosity
seeds" — unresolved questions, recurring patterns, or emerging themes that
warrant deeper attention in future interactions.

A curiosity seed is:
  - An open question that was never fully answered
  - A topic or entity that appeared repeatedly but wasn't explored
  - A contradiction or tension that deserves follow-up
  - A pattern that might reveal something important over time

Respond with JSON only (no markdown, no prose):
{
  "seeds": [
    {
      "question": "<open question or pattern in ≤20 words>",
      "theme": "<short theme label, e.g. 'work_habits', 'recurring_entity'>",
      "source_claim_ids": ["<id1>", "<id2>"],
      "confidence": 0.0-1.0
    }
  ]
}
Return 2–5 seeds. If fewer than 2 are genuinely worth tracking, return those only."""

_CURIOSITY_USER = """\
Session memories to analyze (last {n} claims):

{claims_json}

Identify curiosity seeds from these memories."""

_THEME_SYSTEM = """\
You are a thematic clustering assistant for a memory system. Given a list of
memory claims, group them by theme. A theme is a shared entity, topic, or
concept that appears across multiple claims.

Respond with JSON only:
{
  "themes": [
    {
      "label": "<short theme name>",
      "claim_ids": ["<id1>", "<id2>"],
      "summary": "<1-sentence description of the theme>"
    }
  ]
}
Produce 1–8 themes. Only group claims that genuinely share a theme. Skip singletons."""

_THEME_USER = """\
Memory claims to cluster:

{claims_json}

Identify thematic clusters."""

_JOURNAL_SYSTEM = """\
You are a dream journal writer for an AI memory system. You have just completed
a full dream cycle consolidation. Write a brief, human-readable markdown summary
of what was processed.

Be direct and factual. Use markdown headers and bullet points. Keep it under 300 words."""

_JOURNAL_USER = """\
Dream cycle statistics:
{stats_json}

Curiosity seeds planted:
{seeds_json}

Thematic clusters discovered:
{themes_json}

Cross-memory edges created: {edges_count}

Write the dream journal entry."""


# ---------------------------------------------------------------------------
# Report dataclasses
# ---------------------------------------------------------------------------

@dataclass
class CuriositySeed:
    """A single curiosity seed extracted during dream consolidation."""
    id: str = field(default_factory=lambda: f"seed-{uuid.uuid4().hex[:12]}")
    question: str = ""
    theme: str = ""
    source_claim_ids: List[str] = field(default_factory=list)
    confidence: float = 0.5
    owner_id: str = "default"
    created_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )


@dataclass
class ThemeCluster:
    """A thematic cluster of related memory claims."""
    label: str = ""
    claim_ids: List[str] = field(default_factory=list)
    summary: str = ""


@dataclass
class CuriosityReport:
    """Full report from one CuriosityEngine.run() call."""
    owner_id: str
    seeds: List[CuriositySeed] = field(default_factory=list)
    themes: List[ThemeCluster] = field(default_factory=list)
    edges_created: int = 0
    journal_md: str = ""
    duration_ms: int = 0
    completed_at: str = field(
        default_factory=lambda: datetime.now(timezone.utc).isoformat()
    )

    def to_dict(self) -> dict:
        return {
            "owner_id": self.owner_id,
            "seeds": [
                {
                    "id": s.id,
                    "question": s.question,
                    "theme": s.theme,
                    "source_claim_ids": s.source_claim_ids,
                    "confidence": s.confidence,
                }
                for s in self.seeds
            ],
            "themes": [
                {
                    "label": t.label,
                    "claim_ids": t.claim_ids,
                    "summary": t.summary,
                }
                for t in self.themes
            ],
            "edges_created": self.edges_created,
            "duration_ms": self.duration_ms,
            "completed_at": self.completed_at,
        }


# ---------------------------------------------------------------------------
# CuriosityEngine
# ---------------------------------------------------------------------------

class CuriosityEngine:
    """
    Implements M6c: Curiosity Seeds + Cross-Memory Connections.

    Usage (from SleepConsolidator.dream()):
        engine = CuriosityEngine(storage, llm, config)
        report = await engine.run(
            owner_id="default",
            session_claims=consolidated_claims,
            dream_stats={...},
        )
        dream_report.journal = report.journal_md
        # seeds already stored as SemanticClaims by engine.run()
    """

    def __init__(
        self,
        storage: "StorageAdapter",
        llm: "LLMProvider",
        config: Optional["CortexConfig"] = None,
    ) -> None:
        self.storage = storage
        self.llm = llm
        self.config = config
        # Config knobs (with defaults)
        self._max_claims_for_curiosity = getattr(config, "curiosity_max_claims", 50)
        self._min_theme_size = getattr(config, "curiosity_min_theme_size", 2)
        self._store_seeds = getattr(config, "curiosity_store_seeds", True)

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    async def run(
        self,
        owner_id: str,
        session_claims: List["SemanticClaim"],
        dream_stats: Optional[dict] = None,
    ) -> CuriosityReport:
        """
        Full M6c pipeline:
          1. Extract curiosity seeds via LLM
          2. Cluster session claims by theme via LLM
          3. Create cross-memory edges for theme clusters
          4. Store seeds as SemanticClaims (memory_type='curiosity_seed')
          5. Generate human-readable dream journal

        Returns CuriosityReport with all results.
        """
        t0 = time.monotonic()
        report = CuriosityReport(owner_id=owner_id)

        if not session_claims:
            logger.info("[M6c] no session claims — skipping curiosity extraction")
            report.journal_md = _minimal_journal(dream_stats or {})
            report.duration_ms = int((time.monotonic() - t0) * 1000)
            return report

        # Truncate to avoid token blow-out on very large sessions
        claims_sample = session_claims[-self._max_claims_for_curiosity :]

        # Step 1: Curiosity seeds
        seeds = await self._extract_curiosity_seeds(owner_id, claims_sample)
        report.seeds = seeds

        # Step 2: Thematic clustering
        themes = await self._cluster_by_theme(claims_sample)
        report.themes = themes

        # Step 3: Cross-memory edges
        edges_created = await self._create_cross_memory_edges(
            owner_id, themes, claims_sample
        )
        report.edges_created = edges_created

        # Step 4: Persist seeds as special memory claims
        if self._store_seeds and seeds:
            await self._store_curiosity_seeds(owner_id, seeds, session_claims)

        # Step 5: Generate journal
        journal_md = await self._generate_journal(
            seeds, themes, edges_created, dream_stats or {}
        )
        report.journal_md = journal_md

        report.duration_ms = int((time.monotonic() - t0) * 1000)
        report.completed_at = datetime.now(timezone.utc).isoformat()
        logger.info(
            "[M6c] done — seeds=%d themes=%d edges=%d duration_ms=%d",
            len(seeds), len(themes), edges_created, report.duration_ms,
        )
        return report

    # ------------------------------------------------------------------
    # Step 1: Curiosity seed extraction
    # ------------------------------------------------------------------

    async def _extract_curiosity_seeds(
        self, owner_id: str, claims: List["SemanticClaim"]
    ) -> List[CuriositySeed]:
        """Ask the LLM to identify 2-5 curiosity seeds from session claims."""
        claims_json = _claims_to_json(claims)
        user_msg = _CURIOSITY_USER.format(
            n=len(claims), claims_json=claims_json
        )
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "curiosity_seeds")
            _resp = await self.llm.complete(
                system=_CURIOSITY_SYSTEM,
                user=user_msg,
                **({"model": _routed_model} if _routed_model else {}),
            )
            raw = _resp.text if hasattr(_resp, "text") else str(_resp)
        except Exception as exc:
            logger.warning("[M6c] LLM call for curiosity seeds failed: %s", exc)
            return []

        seeds: List[CuriositySeed] = []
        try:
            data = _parse_json(raw)
            for item in data.get("seeds", []):
                q = item.get("question", "").strip()
                if not q:
                    continue
                seed = CuriositySeed(
                    question=q,
                    theme=item.get("theme", ""),
                    source_claim_ids=item.get("source_claim_ids", []),
                    confidence=float(item.get("confidence", 0.5)),
                    owner_id=owner_id,
                )
                seeds.append(seed)
                if len(seeds) >= 5:
                    break
        except Exception as exc:
            logger.warning("[M6c] failed to parse curiosity seed response: %s", exc)

        logger.debug("[M6c] curiosity seeds extracted: %d", len(seeds))
        return seeds

    # ------------------------------------------------------------------
    # Step 2: Thematic clustering
    # ------------------------------------------------------------------

    async def _cluster_by_theme(
        self, claims: List["SemanticClaim"]
    ) -> List[ThemeCluster]:
        """Ask the LLM to cluster claims by theme."""
        claims_json = _claims_to_json(claims)
        user_msg = _THEME_USER.format(claims_json=claims_json)
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "curiosity_themes")
            _resp = await self.llm.complete(
                system=_THEME_SYSTEM,
                user=user_msg,
                **({"model": _routed_model} if _routed_model else {}),
            )
            raw = _resp.text if hasattr(_resp, "text") else str(_resp)
        except Exception as exc:
            logger.warning("[M6c] LLM call for thematic clustering failed: %s", exc)
            return []

        themes: List[ThemeCluster] = []
        try:
            data = _parse_json(raw)
            for item in data.get("themes", []):
                label = item.get("label", "").strip()
                claim_ids = item.get("claim_ids", [])
                if not label or len(claim_ids) < self._min_theme_size:
                    continue
                themes.append(
                    ThemeCluster(
                        label=label,
                        claim_ids=claim_ids,
                        summary=item.get("summary", ""),
                    )
                )
        except Exception as exc:
            logger.warning("[M6c] failed to parse theme clusters: %s", exc)

        logger.debug("[M6c] theme clusters: %d", len(themes))
        return themes

    # ------------------------------------------------------------------
    # Step 3: Cross-memory edge creation
    # ------------------------------------------------------------------

    async def _create_cross_memory_edges(
        self,
        owner_id: str,
        themes: List[ThemeCluster],
        claims: List["SemanticClaim"],
    ) -> int:
        """
        For each theme cluster, create RELATED_TO edges between all pairs
        of claims that share that theme (up to max_pairs to avoid n² explosion).
        Returns count of edges successfully created.
        """
        MAX_PAIRS_PER_THEME = 10
        edges_created = 0
        # Build a quick id → claim lookup
        claim_map: Dict[str, "SemanticClaim"] = {c.id: c for c in claims}

        for theme in themes:
            valid_ids = [cid for cid in theme.claim_ids if cid in claim_map]
            if len(valid_ids) < 2:
                continue

            # Create edges for all pairs (up to limit)
            pairs: List[Tuple[str, str]] = []
            for i in range(len(valid_ids)):
                for j in range(i + 1, len(valid_ids)):
                    pairs.append((valid_ids[i], valid_ids[j]))
                    if len(pairs) >= MAX_PAIRS_PER_THEME:
                        break
                if len(pairs) >= MAX_PAIRS_PER_THEME:
                    break

            for src_id, tgt_id in pairs:
                try:
                    await self.storage.create_edge(
                        source_type="claim",
                        source_id=src_id,
                        target_type="claim",
                        target_id=tgt_id,
                        edge_type="RELATED_TO",
                        owner_id=owner_id,
                        metadata_json=json.dumps(
                            {
                                "theme": theme.label,
                                "theme_summary": theme.summary,
                                "created_by": "M6c_curiosity",
                            }
                        ),
                    )
                    edges_created += 1
                except Exception as exc:
                    logger.debug(
                        "[M6c] edge creation skipped (%s → %s): %s",
                        src_id, tgt_id, exc,
                    )

        logger.debug("[M6c] cross-memory edges created: %d", edges_created)
        return edges_created

    # ------------------------------------------------------------------
    # Step 4: Store curiosity seeds as SemanticClaims
    # ------------------------------------------------------------------

    async def _store_curiosity_seeds(
        self,
        owner_id: str,
        seeds: List[CuriositySeed],
        all_claims: List["SemanticClaim"],
    ) -> None:
        """
        Store each seed as a SemanticClaim with memory_type='curiosity_seed'.
        This makes seeds visible at next wake() as pending curiosities.
        """
        # Find a session_id to attach seeds to (use the most recent claim's session)
        source_session = None
        if all_claims:
            source_session = all_claims[-1].source_session

        for seed in seeds:
            try:
                await self.storage.create_claim(
                    owner_id=owner_id,
                    subject="curiosity_seed",
                    predicate="asks",
                    object_value=seed.question,
                    memory_type="curiosity_seed",
                    confidence=seed.confidence,
                    status="active",
                    source_session=source_session,
                    metadata_json=json.dumps(
                        {
                            "seed_id": seed.id,
                            "theme": seed.theme,
                            "source_claim_ids": seed.source_claim_ids,
                            "created_by": "M6c_curiosity",
                        }
                    ),
                )
                logger.debug("[M6c] stored seed: %s (theme=%s)", seed.id, seed.theme)
            except Exception as exc:
                logger.warning(
                    "[M6c] failed to store seed %s: %s", seed.id, exc
                )

    # ------------------------------------------------------------------
    # Step 5: Dream journal generation
    # ------------------------------------------------------------------

    async def _generate_journal(
        self,
        seeds: List[CuriositySeed],
        themes: List[ThemeCluster],
        edges_created: int,
        dream_stats: dict,
    ) -> str:
        """Generate a human-readable markdown dream journal via LLM."""
        seeds_json = json.dumps(
            [{"question": s.question, "theme": s.theme} for s in seeds],
            ensure_ascii=False,
            indent=2,
        )
        themes_json = json.dumps(
            [{"label": t.label, "summary": t.summary, "size": len(t.claim_ids)} for t in themes],
            ensure_ascii=False,
            indent=2,
        )
        stats_json = json.dumps(dream_stats, ensure_ascii=False, indent=2)

        user_msg = _JOURNAL_USER.format(
            stats_json=stats_json,
            seeds_json=seeds_json,
            themes_json=themes_json,
            edges_count=edges_created,
        )
        try:
            from cortex.llm.router import resolve_routed_model
            _routed_model = resolve_routed_model(self.llm, "curiosity_journal")
            _resp = await self.llm.complete(
                system=_JOURNAL_SYSTEM,
                user=user_msg,
                **({"model": _routed_model} if _routed_model else {}),
            )
            journal = _resp.text if hasattr(_resp, "text") else str(_resp)
            return journal.strip()
        except Exception as exc:
            logger.warning("[M6c] journal generation failed: %s", exc)
            # Fallback to template journal
            return _fallback_journal(seeds, themes, edges_created, dream_stats)


# ---------------------------------------------------------------------------
# Internal helpers
# ---------------------------------------------------------------------------

def _claims_to_json(claims: List["SemanticClaim"]) -> str:
    """Compact JSON representation of claims for LLM prompts."""
    rows = []
    for c in claims:
        rows.append(
            {
                "id": c.id,
                "subject": c.subject,
                "predicate": c.predicate,
                "object": c.object_value,
                "memory_type": c.memory_type,
                "confidence": round(c.confidence, 2),
            }
        )
    return json.dumps(rows, ensure_ascii=False, indent=2)


def _parse_json(text: str) -> dict:
    """Safely parse LLM JSON output, stripping markdown fences if present."""
    text = text.strip()
    # Strip ```json ... ``` or ``` ... ``` fences
    if text.startswith("```"):
        lines = text.splitlines()
        # Remove first and last lines if they're fences
        inner = []
        in_block = False
        for line in lines:
            if line.startswith("```") and not in_block:
                in_block = True
                continue
            if line.startswith("```") and in_block:
                break
            if in_block:
                inner.append(line)
        text = "\n".join(inner).strip()
    try:
        data = json.loads(text)
        if isinstance(data, dict):
            return data
        return {}
    except json.JSONDecodeError:
        logger.warning("[M6c] JSON parse failed: %r", text[:200])
        return {}


def _minimal_journal(stats: dict) -> str:
    """Fallback journal when there are no session claims."""
    return (
        "# Dream Journal\n\n"
        "_No session memories to process — dream cycle ran on an empty session._\n\n"
        f"**Stats:** {json.dumps(stats)}\n"
    )


def _fallback_journal(
    seeds: List[CuriositySeed],
    themes: List[ThemeCluster],
    edges_created: int,
    stats: dict,
) -> str:
    """Template-based journal when LLM call fails."""
    lines = ["# Dream Journal\n"]

    if stats:
        lines.append("## Cycle Statistics")
        for k, v in stats.items():
            lines.append(f"- **{k}:** {v}")
        lines.append("")

    if seeds:
        lines.append("## Curiosity Seeds Planted")
        for s in seeds:
            lines.append(f"- [{s.theme}] {s.question} _(confidence: {s.confidence:.2f})_")
        lines.append("")

    if themes:
        lines.append("## Themes Discovered")
        for t in themes:
            lines.append(f"- **{t.label}** ({len(t.claim_ids)} claims) — {t.summary}")
        lines.append("")

    lines.append(f"## Cross-Memory Connections\n- {edges_created} RELATED_TO edges created\n")

    return "\n".join(lines)
