#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 9: Pheromone Replayer & Drift Analyzer

Given a pheromone JSONL file, this tool can:

- Filter entries by task_type / gate / helix_program
- Re-run those tasks against the current BloomOS/Helix64 runtime
- Compare old vs new:
    - status (ok/error)
    - logits_mean
    - Se (symbolic entropy)
    - fitness score

This is useful for:
- Regression testing (did a substrate change break behavior?)
- Sanity checks on adaptive controllers
- Measuring drift in performance over time

Usage:
    # Dry-run: just list original metrics
    PYTHONPATH=. python3 -m helix_cdc.bloomos.pheromone_replay \\
      /tmp/pheromones.jsonl --task-type math --limit 5 --dry-run

    # Full replay against live BloomOS
    PYTHONPATH=. python3 -m helix_cdc.bloomos.pheromone_replay \\
      /tmp/pheromones.jsonl --task-type code --limit 5
"""

import json
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Dict, Iterator, List, Optional, Tuple

from .bloomos_client import BloomOSClient, TaskResult
from .pheromone import PheromoneBoard, PheromoneEntry
from .fitness import fitness
from helix_cdc.metrics import measure_text


@dataclass
class ReplayConfig:
    """Configuration for pheromone replay."""
    path: Path
    task_type: Optional[str] = None
    gate: Optional[str] = None
    helix_program: Optional[str] = None
    limit: Optional[int] = None
    base_url: str = "http://localhost:7070"
    timeout: float = 60.0
    dry_run: bool = False


@dataclass
class ReplayResult:
    """Result of replaying a single pheromone entry."""
    # Identity
    original_index: int
    task_type: str
    gate: Optional[str]
    helix_program: Optional[str]
    goal: Optional[str]

    # Original metrics
    original_status: str
    original_logits_mean: Optional[float]
    original_Se: Optional[float]
    original_fitness: Optional[float]

    # New run metrics
    new_status: str
    new_logits_mean: Optional[float]
    new_Se: Optional[float]
    new_fitness: Optional[float]

    # Drift (new - original)
    logits_drift: Optional[float]
    Se_drift: Optional[float]
    fitness_drift: Optional[float]


@dataclass
class ReplaySummary:
    """Summary statistics across all replays."""
    total_replayed: int
    total_skipped: int
    status_matches: int
    status_mismatches: int
    avg_logits_drift: Optional[float]
    avg_Se_drift: Optional[float]
    avg_fitness_drift: Optional[float]
    regressions: List[int]  # indices where new fitness < original


# -------------------------------------------------------------------
# Helpers
# -------------------------------------------------------------------

def _iter_filtered_entries(cfg: ReplayConfig) -> Iterator[Tuple[int, PheromoneEntry]]:
    """Iterate over pheromone entries matching the filter criteria."""
    board = PheromoneBoard(str(cfg.path))
    count = 0
    for i, entry in enumerate(board.iter_entries()):
        if cfg.task_type and entry.task_type != cfg.task_type:
            continue
        if cfg.gate and entry.gate != cfg.gate:
            continue
        if cfg.helix_program and entry.helix_program != cfg.helix_program:
            continue
        yield i, entry
        count += 1
        if cfg.limit is not None and count >= cfg.limit:
            break


def _extract_goal(entry: PheromoneEntry) -> Optional[str]:
    """Extract the original goal from pheromone entry meta."""
    if entry.meta and isinstance(entry.meta, dict):
        goal = entry.meta.get("goal")
        if isinstance(goal, str):
            return goal
    return None


def _compute_Se_from_text(text: Optional[str]) -> Optional[float]:
    """Compute Se from output text."""
    if not text or not text.strip():
        return None
    try:
        m = measure_text(text, auto_detect_d=True)
        return m.Se
    except Exception:
        return None


def _extract_text_from_result(res: TaskResult) -> Optional[str]:
    """Extract text from TaskResult output."""
    out = res.output
    if out is None:
        return None
    if isinstance(out, str):
        return out
    if isinstance(out, dict):
        for key in ("text", "content", "response", "output", "result", "reasoning"):
            if key in out and isinstance(out[key], str):
                return out[key]
        return json.dumps(out)
    return str(out)


class _FakeEntry:
    """Helper class to compute fitness from dict-like data."""
    def __init__(self, d: Dict[str, Any]) -> None:
        self.__dict__.update(d)


def _run_single_replay(client: BloomOSClient, idx: int, entry: PheromoneEntry) -> ReplayResult:
    """Replay a single pheromone entry against the runtime."""
    goal = _extract_goal(entry)
    task_type = entry.task_type or "unknown"
    gate = entry.gate
    helix_program = entry.helix_program

    # Original metrics
    orig_status = entry.status or "unknown"
    orig_logits = entry.logits_mean
    orig_Se = entry.se_Se
    try:
        orig_fitness = fitness(entry)
    except Exception:
        orig_fitness = None

    # If we don't have a goal, we can't faithfully re-run the task.
    if not goal:
        return ReplayResult(
            original_index=idx,
            task_type=task_type,
            gate=gate,
            helix_program=helix_program,
            goal=None,
            original_status=orig_status,
            original_logits_mean=orig_logits,
            original_Se=orig_Se,
            original_fitness=orig_fitness,
            new_status="skipped_no_goal",
            new_logits_mean=None,
            new_Se=None,
            new_fitness=None,
            logits_drift=None,
            Se_drift=None,
            fitness_drift=None,
        )

    # Build extra payload
    extra: Dict[str, Any] = {}
    if helix_program:
        extra["helix_program"] = helix_program

    # Re-run task
    try:
        res = client.run_task(
            task_type=task_type,
            description=goal,
            anchors=None,
            gate_policy=gate,
            extra=extra if extra else None,
        )
        new_status = res.status
        new_logits = res.logits_mean
        text_out = _extract_text_from_result(res)
        new_Se = _compute_Se_from_text(text_out)
        logits_std = res.logits_std
    except Exception as e:
        # Runtime error
        return ReplayResult(
            original_index=idx,
            task_type=task_type,
            gate=gate,
            helix_program=helix_program,
            goal=goal,
            original_status=orig_status,
            original_logits_mean=orig_logits,
            original_Se=orig_Se,
            original_fitness=orig_fitness,
            new_status=f"error:{type(e).__name__}",
            new_logits_mean=None,
            new_Se=None,
            new_fitness=None,
            logits_drift=None,
            Se_drift=None,
            fitness_drift=None,
        )

    # Compute new fitness
    fake_entry = _FakeEntry({
        "task_type": task_type,
        "gate": gate,
        "helix_program": helix_program,
        "status": new_status,
        "logits_mean": new_logits,
        "logits_std": logits_std,
        "se_H": None,
        "se_C": None,
        "se_D": None,
        "se_Se": new_Se,
        "seed_len": None,
        "seed_entropy": entry.seed_entropy,  # Use original seed_entropy for comparison
        "regrow_depth": None,
        "lattice_D": None,
        "meta": {},
    })
    try:
        new_fitness = fitness(fake_entry)
    except Exception:
        new_fitness = None

    # Drift calculations
    logits_drift = None
    if orig_logits is not None and new_logits is not None:
        logits_drift = new_logits - orig_logits

    Se_drift = None
    if orig_Se is not None and new_Se is not None:
        Se_drift = new_Se - orig_Se

    fitness_drift = None
    if orig_fitness is not None and new_fitness is not None:
        fitness_drift = new_fitness - orig_fitness

    return ReplayResult(
        original_index=idx,
        task_type=task_type,
        gate=gate,
        helix_program=helix_program,
        goal=goal,
        original_status=orig_status,
        original_logits_mean=orig_logits,
        original_Se=orig_Se,
        original_fitness=orig_fitness,
        new_status=new_status,
        new_logits_mean=new_logits,
        new_Se=new_Se,
        new_fitness=new_fitness,
        logits_drift=logits_drift,
        Se_drift=Se_drift,
        fitness_drift=fitness_drift,
    )


def _compute_summary(results: List[ReplayResult]) -> ReplaySummary:
    """Compute summary statistics from replay results."""
    total_replayed = 0
    total_skipped = 0
    status_matches = 0
    status_mismatches = 0
    logits_drifts: List[float] = []
    Se_drifts: List[float] = []
    fitness_drifts: List[float] = []
    regressions: List[int] = []

    for r in results:
        if r.new_status.startswith("skipped") or r.new_status.startswith("dry_run"):
            total_skipped += 1
            continue

        total_replayed += 1

        # Status comparison (normalize to ok/error)
        orig_ok = r.original_status == "ok"
        new_ok = r.new_status == "ok"
        if orig_ok == new_ok:
            status_matches += 1
        else:
            status_mismatches += 1

        if r.logits_drift is not None:
            logits_drifts.append(r.logits_drift)
        if r.Se_drift is not None:
            Se_drifts.append(r.Se_drift)
        if r.fitness_drift is not None:
            fitness_drifts.append(r.fitness_drift)
            if r.fitness_drift < 0:
                regressions.append(r.original_index)

    return ReplaySummary(
        total_replayed=total_replayed,
        total_skipped=total_skipped,
        status_matches=status_matches,
        status_mismatches=status_mismatches,
        avg_logits_drift=sum(logits_drifts) / len(logits_drifts) if logits_drifts else None,
        avg_Se_drift=sum(Se_drifts) / len(Se_drifts) if Se_drifts else None,
        avg_fitness_drift=sum(fitness_drifts) / len(fitness_drifts) if fitness_drifts else None,
        regressions=regressions,
    )


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

def replay(cfg: ReplayConfig) -> List[ReplayResult]:
    """
    Replay pheromone entries against the current runtime.

    Args:
        cfg: ReplayConfig with path, filters, and runtime settings

    Returns:
        List of ReplayResult for each entry
    """
    client = BloomOSClient(base_url=cfg.base_url, timeout=cfg.timeout)

    results: List[ReplayResult] = []
    for idx, entry in _iter_filtered_entries(cfg):
        if cfg.dry_run:
            # Only record original metrics, don't re-run
            try:
                orig_fitness = fitness(entry)
            except Exception:
                orig_fitness = None
            results.append(
                ReplayResult(
                    original_index=idx,
                    task_type=entry.task_type or "unknown",
                    gate=entry.gate,
                    helix_program=entry.helix_program,
                    goal=_extract_goal(entry),
                    original_status=entry.status or "unknown",
                    original_logits_mean=entry.logits_mean,
                    original_Se=entry.se_Se,
                    original_fitness=orig_fitness,
                    new_status="dry_run",
                    new_logits_mean=None,
                    new_Se=None,
                    new_fitness=None,
                    logits_drift=None,
                    Se_drift=None,
                    fitness_drift=None,
                )
            )
        else:
            rr = _run_single_replay(client, idx, entry)
            results.append(rr)
    return results


def replay_with_summary(cfg: ReplayConfig) -> Tuple[List[ReplayResult], ReplaySummary]:
    """
    Replay pheromone entries and compute summary statistics.

    Returns:
        Tuple of (results list, summary)
    """
    results = replay(cfg)
    summary = _compute_summary(results)
    return results, summary


# -------------------------------------------------------------------
# CLI
# -------------------------------------------------------------------

def _cli() -> None:
    import argparse

    parser = argparse.ArgumentParser(
        prog="pheromone_replay",
        description="Phase 9: Replay pheromone entries against current BloomOS/Helix64 runtime.",
    )
    parser.add_argument(
        "path",
        help="Path to pheromone JSONL file.",
    )
    parser.add_argument(
        "--task-type",
        help="Filter by task_type.",
    )
    parser.add_argument(
        "--gate",
        help="Filter by gate policy.",
    )
    parser.add_argument(
        "--helix-program",
        help="Filter by helix_program sequence.",
    )
    parser.add_argument(
        "--limit",
        type=int,
        help="Maximum number of entries to replay.",
    )
    parser.add_argument(
        "--base-url",
        default="http://localhost:7070",
        help="BloomOS runtime base URL (default: http://localhost:7070).",
    )
    parser.add_argument(
        "--timeout",
        type=float,
        default=60.0,
        help="HTTP timeout in seconds (default: 60.0).",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Do not re-run tasks, only report original metrics.",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output full JSON instead of a human-readable table.",
    )
    parser.add_argument(
        "--summary",
        action="store_true",
        help="Show summary statistics at the end.",
    )

    args = parser.parse_args()

    cfg = ReplayConfig(
        path=Path(args.path),
        task_type=args.task_type,
        gate=args.gate,
        helix_program=args.helix_program,
        limit=args.limit,
        base_url=args.base_url,
        timeout=args.timeout,
        dry_run=args.dry_run,
    )

    results, summary = replay_with_summary(cfg)

    if args.json:
        output = {
            "results": [asdict(r) for r in results],
            "summary": asdict(summary),
        }
        print(json.dumps(output, indent=2))
        return

    # Human-readable view
    print("=" * 80)
    print("PHEROMONE REPLAY RESULTS")
    print("=" * 80)
    print(
        f"File: {cfg.path}\n"
        f"Filters: task_type={cfg.task_type or '*'} | "
        f"gate={cfg.gate or '*'} | "
        f"program={cfg.helix_program or '*'}\n"
        f"Mode: {'DRY RUN' if cfg.dry_run else 'LIVE REPLAY'}"
    )
    print()

    header = (
        "idx  task    gate        program     "
        "orig_st  new_st   d_logits   d_Se    d_fit"
    )
    print(header)
    print("-" * len(header))

    for r in results:
        d_logits = f"{r.logits_drift:+.4f}" if r.logits_drift is not None else "   -   "
        d_Se = f"{r.Se_drift:+.2f}" if r.Se_drift is not None else "  -  "
        d_fit = f"{r.fitness_drift:+.3f}" if r.fitness_drift is not None else "  -  "
        gate = (r.gate or "-")[:10]
        prog = (r.helix_program or "-")[:10]
        orig_st = r.original_status[:7]
        new_st = r.new_status[:7]
        print(
            f"{r.original_index:<4d} "
            f"{r.task_type[:5]:<5} "
            f"{gate:<10} "
            f"{prog:<10} "
            f"{orig_st:<7} "
            f"{new_st:<7} "
            f"{d_logits:>8} "
            f"{d_Se:>6} "
            f"{d_fit:>7}"
        )

    if args.summary or not cfg.dry_run:
        print()
        print("=" * 80)
        print("SUMMARY")
        print("=" * 80)
        print(f"Total replayed:      {summary.total_replayed}")
        print(f"Total skipped:       {summary.total_skipped}")
        print(f"Status matches:      {summary.status_matches}")
        print(f"Status mismatches:   {summary.status_mismatches}")
        if summary.avg_logits_drift is not None:
            print(f"Avg logits drift:    {summary.avg_logits_drift:+.4f}")
        if summary.avg_Se_drift is not None:
            print(f"Avg Se drift:        {summary.avg_Se_drift:+.2f}")
        if summary.avg_fitness_drift is not None:
            print(f"Avg fitness drift:   {summary.avg_fitness_drift:+.3f}")
        if summary.regressions:
            print(f"Regressions (indices): {summary.regressions}")
        else:
            print("Regressions:         None detected")


if __name__ == "__main__":
    _cli()
