# echo_cog/runtime.py
"""
Echo Cog Runtime - Wraps a GGUF model with Se state and FibPi3D wave scheduling.

This is the main interface for Echo Cog:
  - Loads model via llama-cli
  - Maintains SeState across questions
  - Applies FibPi3D wave routing
  - Builds prompts with task + Se summary + wave info
"""

from __future__ import annotations

import json
import subprocess
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import List, Literal, Optional, Tuple

from .se_state import SeState
from .fibpi3d_scheduler import FibPi3DScheduler, WaveInfo
from .backend_llama_cpp import LlamaCppBackendConfig, run_llama_cpp
from .se_controller_adapter import SeControllerAdapter, CogSeDecision, get_se_adapter, is_se_controller_available


# Profile types for Echo Cog runtime
ProfileName = Literal["proof", "fast_mmlu", "chat"]

# Mode types for cognitive behavior
ModeName = Literal["minimal", "rich"]

# Backend types for inference engine
BackendName = Literal["llama_cli", "llama_cpp"]

# Prompt style types (for isolating prompt-format effects)
PromptStyle = Literal["substrate_simple", "cog_minimal", "cog_rich", "regen_safe"]

# Rich mode system prompt (from DESIGN_NOTES_70PCT)
RICH_SYSTEM_PROMPT = """You are Echo Cog, a careful reasoning assistant.
For each question, you will first think step by step in [THOUGHTS],
then give the final answer as a single letter (A, B, C, or D) in [ANSWER].
Never explain outside [THOUGHTS]. [ANSWER] must contain only one letter."""


@dataclass
class MMLUSample:
    """A single MMLU question."""
    task: str
    question: str
    choices: List[str]
    answer: str  # "A"/"B"/"C"/"D"


def load_mmlu_file(path: Path, task_name: Optional[str] = None) -> List[MMLUSample]:
    """Load MMLU samples from a .jsonl file."""
    samples: List[MMLUSample] = []
    task = task_name or path.stem
    with path.open("r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            row = json.loads(line)
            q = row.get("question") or row.get("q")
            if not q:
                raise ValueError(f"Missing question/q in {path}")
            choices = row["choices"]
            ans = row["answer"]
            if isinstance(ans, int):
                # integer index → letter
                if ans < 0 or ans >= len(choices):
                    raise ValueError(f"Bad answer index {ans} in {path}")
                letter = chr(ord("A") + ans)
            else:
                letter = str(ans).strip().upper()
            samples.append(
                MMLUSample(
                    task=task,
                    question=q,
                    choices=choices,
                    answer=letter,
                )
            )
    return samples


def call_llama_cli(
    llama_cli: Path,
    model: Path,
    prompt: str,
    max_tokens: int = 4,
    timeout: int = 120,
    threads: int = 8,
    ctx_size: int = 2048,
    extra_args: Optional[List[str]] = None,
) -> str:
    """Call llama-cli and return the output."""
    args = [
        str(llama_cli),
        "-m",
        str(model),
        "-p",
        prompt,
        "-n",
        str(max_tokens),
        "--top-k",
        "1",  # greedy-ish
        "-no-cnv",
        "--threads",
        str(threads),
        "--ctx-size",
        str(ctx_size),
    ]
    if extra_args:
        args.extend(extra_args)
    proc = subprocess.run(
        args,
        stdout=subprocess.PIPE,
        stderr=subprocess.STDOUT,
        timeout=timeout,
    )
    return proc.stdout.decode("utf-8", errors="replace")


def extract_letter(output: str, prefer_answer_block: bool = True) -> Optional[str]:
    """
    Extract the answer letter from model output.

    If prefer_answer_block is True, first look for a letter near [ANSWER],
    then fall back to the last A/B/C/D in the output.
    """
    import re

    # Strategy 1: Look for letter after [ANSWER] marker
    if prefer_answer_block:
        # Find [ANSWER] section and extract letter from there
        answer_match = re.search(r'\[ANSWER\][:\s]*([A-Da-d])', output, re.IGNORECASE)
        if answer_match:
            return answer_match.group(1).upper()

    # Strategy 2: Look for "Answer:" or "answer:" pattern
    direct_match = re.search(r'[Aa]nswer[:\s]+([A-Da-d])', output)
    if direct_match:
        return direct_match.group(1).upper()

    # Strategy 3: Fall back to last A/B/C/D in output
    for ch in reversed(output):
        up = ch.upper()
        if up in ("A", "B", "C", "D"):
            return up
    return None


@dataclass
class EchoCogConfig:
    """Configuration for an Echo Cog runtime instance."""
    llama_cli: Path
    model_path: Path
    system_tag: str = "Echo Cog Mistral"
    max_tokens: int = 4
    timeout: int = 120
    num_ports: int = 8
    profile: ProfileName = "proof"
    threads: int = 8
    ctx_size: int = 2048
    extra_llama_args: List[str] = field(default_factory=list)
    mode: ModeName = "minimal"  # "minimal" or "rich" cognitive mode
    history_k: int = 5  # Number of Q/A pairs to include in rich mode history
    backend: BackendName = "llama_cli"  # "llama_cli" (subprocess) or "llama_cpp" (warm model)
    use_se_controller: bool = False  # Enable SeController-driven behavior
    se_topology: str = "fibpi"  # Topology for Se decisions: "fibpi", "shell", "1d"
    prompt_style: PromptStyle = "cog_minimal"  # Prompt format: substrate_simple, cog_minimal, cog_rich, regen_safe


class EchoCogRuntime:
    """
    Echo Cog wrapper around a single GGUF model (regen or vendor).

    Responsibilities:
      - Maintain SeState across questions
      - Request wave routing from FibPi3DScheduler
      - Build prompts that include task + Se summary + wave info
      - Call llama-cli and parse a single-letter answer
    """

    def __init__(
        self,
        config: EchoCogConfig,
        se_state: Optional[SeState] = None,
        scheduler: Optional[FibPi3DScheduler] = None,
    ) -> None:
        self.config = config
        self.se = se_state or SeState()
        self.scheduler = scheduler or FibPi3DScheduler(num_ports=config.num_ports)

        # Set up llama_cpp backend config if using warm model
        if self.config.backend == "llama_cpp":
            self._llama_cpp_cfg = LlamaCppBackendConfig(
                model_path=str(config.model_path),
                ctx_size=config.ctx_size,
                threads=config.threads,
            )
        else:
            self._llama_cpp_cfg = None

        # Set up SeController adapter if enabled
        if self.config.use_se_controller:
            self._se_adapter = SeControllerAdapter(topology=config.se_topology)
        else:
            self._se_adapter = None

    def build_prompt(self, sample: MMLUSample, wave: WaveInfo) -> str:
        """Build a prompt with Se context and wave routing info."""
        # First check prompt_style override - these bypass profile/mode logic
        if self.config.prompt_style == "substrate_simple":
            return self._build_substrate_simple_prompt(sample)

        if self.config.prompt_style == "regen_safe":
            return self._build_regen_safe_prompt(sample)

        # Standard prompt building (cog_minimal or cog_rich)
        choices_lines = [
            f"{chr(ord('A') + i)}. {c}" for i, c in enumerate(sample.choices)
        ]
        choices_block = "\n".join(choices_lines)

        # Fast MMLU profile: minimal prompt for speed
        if self.config.profile == "fast_mmlu":
            return (
                f"Question: {sample.question}\n\n"
                f"{choices_block}\n\n"
                f"Answer (A/B/C/D):"
            )

        # Rich mode (cog_rich or mode="rich"): use explicit [THOUGHTS] + [ANSWER] sections
        if self.config.prompt_style == "cog_rich" or self.config.mode == "rich":
            return self._build_rich_prompt(sample, wave, choices_block)

        # Minimal mode (cog_minimal, default): current proof/chat behavior
        return self._build_minimal_prompt(sample, wave, choices_block)

    def _build_minimal_prompt(
        self, sample: MMLUSample, wave: WaveInfo, choices_block: str
    ) -> str:
        """Build minimal mode prompt (original proof/chat behavior)."""
        se_summary = self.se.summary_for_task(sample.task)
        return (
            f"[SYSTEM] {self.config.system_tag}\n"
            f"[TASK] You are solving MMLU multiple choice questions in the domain '{sample.task}'.\n"
            f"[ROUTING] FibPi3D wave: id={wave.wave_id}, phase={wave.phase:.6f}, port={wave.port}.\n"
            f"[STATE]\n"
            f"{se_summary}\n\n"
            f"[QUESTION]\n"
            f"{sample.question}\n\n"
            f"[CHOICES]\n"
            f"{choices_block}\n\n"
            f"Answer with a single capital letter: A, B, C, or D.\n"
            f"Final answer:"
        )

    def _build_rich_prompt(
        self, sample: MMLUSample, wave: WaveInfo, choices_block: str
    ) -> str:
        """
        Build rich mode prompt with explicit [THOUGHTS] + [ANSWER] sections.

        Based on DESIGN_NOTES_70PCT: assertive system prompt, history context,
        and explicit reasoning scratchpad.
        """
        # Build history block from recent Q/A pairs
        history_block = ""
        recent = self.se.get_recent_qa(sample.task, k=self.config.history_k)
        if recent:
            history_lines = ["[HISTORY]", "(Recent Q/A for this task to calibrate your reasoning.)"]
            for rec in recent:
                status = "correct" if rec.correct else "wrong"
                q_preview = rec.question.strip().replace("\n", " ")
                if len(q_preview) > 100:
                    q_preview = q_preview[:97] + "..."
                history_lines.append(f"Q: {q_preview}")
                history_lines.append(f"A: {rec.model_answer or '?'} ({status})")
                history_lines.append("---")
            history_block = "\n".join(history_lines) + "\n\n"

        return (
            f"[SYSTEM]\n"
            f"{RICH_SYSTEM_PROMPT}\n\n"
            f"[TASK] {sample.task}\n\n"
            f"{history_block}"
            f"[QUESTION]\n"
            f"{sample.question}\n\n"
            f"Choices:\n"
            f"{choices_block}\n\n"
            f"[THOUGHTS]\n"
            f"Think step by step about the question and the choices.\n"
            f"You may compare the options, eliminate wrong answers, and justify your reasoning.\n\n"
            f"[ANSWER]\n"
            f"Now give ONLY the final answer as a single letter (A, B, C, or D):"
        )

    def _build_substrate_simple_prompt(self, sample: MMLUSample) -> str:
        """
        Build substrate-simple prompt - matches the raw substrate KAT format exactly.

        This bypasses all Cog scaffolding (no [SYSTEM], [ROUTING], [STATE] blocks).
        Used to prove that prompt format causes the accuracy gap, not model weights.

        MUST match kat_mistral_mmlu_llama_cpp.py build_prompt() format exactly.
        """
        letters = ["A", "B", "C", "D"]
        opts = "\n".join(f"{l}. {c}" for l, c in zip(letters, sample.choices))

        return f"""You are a multiple-choice question answering system.

Question: {sample.question}

Options:
{opts}

Answer with only the single letter (A, B, C, or D).
Final answer:"""

    def _build_regen_safe_prompt(self, sample: MMLUSample) -> str:
        """
        Build regen-safe prompt - minimal Cog scaffolding that regen handles well.

        Uses simple structure but keeps task context. Falls back to substrate_simple
        style with just a task hint. Useful for debugging regen prompt sensitivity.
        """
        choices_lines = [
            f"{chr(ord('A') + i)}. {c}" for i, c in enumerate(sample.choices)
        ]
        choices_block = "\n".join(choices_lines)

        return (
            f"Task: {sample.task}\n\n"
            f"Question: {sample.question}\n\n"
            f"Choices:\n"
            f"{choices_block}\n\n"
            f"Answer with one letter (A, B, C, or D):"
        )

    def answer_sample(self, sample: MMLUSample) -> Tuple[str, float, str]:
        """
        Answer a single MMLU sample.

        Returns:
            (answer_letter, latency_seconds, raw_output)
        """
        wave = self.scheduler.next_wave(sample.task)

        # Get SeController decision if enabled
        se_decision: Optional[CogSeDecision] = None
        if self._se_adapter is not None:
            try:
                # Get task stats for Se decision
                task_stats = self.se.stats_for_task(sample.task)
                se_decision = self._se_adapter.decide(
                    task=sample.task,
                    task_stats=task_stats,
                    wave_phase=wave.phase,
                    wave_port=wave.port,
                    question=sample.question,
                    choices=sample.choices,
                )
            except Exception as e:
                # Fail-safe: log and continue without Se
                import sys
                print(f"[WARN] SeController.decide failed: {e}", file=sys.stderr)
                se_decision = None

        prompt = self.build_prompt(sample, wave)

        # Determine max_tokens based on profile, mode, and prompt_style
        if self.config.profile == "fast_mmlu":
            max_tokens = 8
        elif self.config.mode == "rich" or self.config.prompt_style == "cog_rich":
            # Rich mode needs more tokens for THOUGHTS section
            max_tokens = 256
        elif self.config.prompt_style == "substrate_simple":
            # EXACTLY match raw substrate KAT (max_tokens=4)
            max_tokens = 4
        elif self.config.prompt_style == "regen_safe":
            # Simple prompts just need enough for a short answer
            max_tokens = 8
        else:
            max_tokens = self.config.max_tokens

        # Let SeController override max_tokens if it has an opinion
        if se_decision is not None and se_decision.max_tokens_hint is not None:
            # Use Se hint, but don't exceed mode's max
            max_tokens = min(max_tokens, se_decision.max_tokens_hint)

        t0 = time.perf_counter()

        # Dispatch to the configured backend
        if self.config.backend == "llama_cpp":
            if self._llama_cpp_cfg is None:
                raise RuntimeError("llama_cpp backend configured but not initialized")
            raw = run_llama_cpp(
                self._llama_cpp_cfg,
                prompt,
                max_tokens=max_tokens,
            )
        else:
            # Default: llama_cli subprocess
            raw = call_llama_cli(
                self.config.llama_cli,
                self.config.model_path,
                prompt,
                max_tokens=max_tokens,
                timeout=self.config.timeout,
                threads=self.config.threads,
                ctx_size=self.config.ctx_size,
                extra_args=self.config.extra_llama_args,
            )

        dt = time.perf_counter() - t0

        # Use answer block extraction for rich mode
        prefer_block = (self.config.mode == "rich")
        letter = extract_letter(raw, prefer_answer_block=prefer_block)
        if not letter:
            letter = "A"  # conservative default

        # Skip Se tracking for fast_mmlu profile (speed optimization)
        if self.config.profile != "fast_mmlu":
            self.se.record(
                task=sample.task,
                question=sample.question,
                choices=sample.choices,
                correct_answer=sample.answer,
                model_answer=letter,
            )

        return letter, dt, raw

    def answer_question(self, question: str, task: str = "general") -> Tuple[str, float]:
        """
        Answer a free-form question (not MMLU format).

        Returns:
            (answer_text, latency_seconds)
        """
        wave = self.scheduler.next_wave(task)
        se_summary = self.se.summary_for_task(task)

        prompt = (
            f"[SYSTEM] {self.config.system_tag}\n"
            f"[TASK] {task}\n"
            f"[ROUTING] FibPi3D wave: id={wave.wave_id}, phase={wave.phase:.6f}, port={wave.port}.\n"
            f"[STATE]\n"
            f"{se_summary}\n\n"
            f"[QUESTION]\n"
            f"{question}\n\n"
            f"Answer:"
        )

        t0 = time.perf_counter()
        raw = call_llama_cli(
            self.config.llama_cli,
            self.config.model_path,
            prompt,
            max_tokens=256,  # More tokens for free-form answers
            timeout=self.config.timeout,
            threads=self.config.threads,
            ctx_size=self.config.ctx_size,
            extra_args=self.config.extra_llama_args,
        )
        dt = time.perf_counter() - t0

        return raw, dt

    def get_stats(self) -> dict:
        """Get current Se state statistics."""
        overall = self.se.overall_stats()
        per_task = {
            task: {"accuracy": stats.accuracy, "asked": stats.asked, "correct": stats.correct}
            for task, stats in self.se.task_stats.items()
        }
        return {
            "overall": overall,
            "per_task": per_task,
            "total_waves": sum(self.scheduler.get_wave_count(t) for t in per_task.keys()),
        }
