#!/usr/bin/env python3
"""Compare two cell-output manifests and run directories to produce a drift report.

Takes manifests generated by ``capture_notebook_outputs.py`` plus the
run-level YAML artefacts (snapshot_grid, project_context, feature profiles)
and produces a structured Markdown drift report.

Usage:
    python scripts/experiments/compare_exploration_runs.py \
        --old-manifest experiments/runs/baseline/cell_outputs.md \
        --new-manifest experiments/runs/current/cell_outputs.md \
        --old-run-dir experiments/runs/baseline \
        --new-run-dir experiments/runs/current \
        --output experiments/drift_report.md
"""

from __future__ import annotations

import argparse
import difflib
import io
import re
import sys
from pathlib import Path
from typing import Any, Optional

try:
    import yaml
except ImportError:
    yaml = None  # type: ignore[assignment]

try:
    import pandas as pd
except ImportError:
    pd = None  # type: ignore[assignment]

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

_DRIFT_THRESHOLD = 0.005
_REGRESSION_THRESHOLD = 0.01

# Regex to split manifest into notebook sections
_NB_HEADER_RE = re.compile(r"^# (\S+)\s*$", re.MULTILINE)

# Regex to split a notebook section into cell blocks
_CELL_HEADER_RE = re.compile(
    r"^## (?:code|config):\s*(\S+)\s+\(id:\s*(\w+)\)", re.MULTILINE
)


# ---------------------------------------------------------------------------
# Manifest parsing
# ---------------------------------------------------------------------------


def _parse_manifest(text: str) -> dict[str, dict[str, str]]:
    """Parse a cell-output manifest into {notebook: {cell_name: content}}.

    Also stores the raw section map under the key ``__section_map__``.
    """
    result: dict[str, dict[str, str]] = {}

    # Extract section map (HTML comments at the top)
    section_map_lines = []
    for line in text.split("\n"):
        if line.startswith("<!--"):
            section_map_lines.append(line)
        elif line.strip() and not line.startswith("<!--"):
            break
    result["__section_map__"] = {"_raw": "\n".join(section_map_lines)}

    # Split by notebook headers
    parts = _NB_HEADER_RE.split(text)
    # parts = ['preamble', 'nb_name', 'nb_content', 'nb_name', 'nb_content', ...]
    i = 1
    while i < len(parts) - 1:
        nb_name = parts[i].strip()
        nb_body = parts[i + 1]
        i += 2

        cells: dict[str, str] = {}
        cell_parts = _CELL_HEADER_RE.split(nb_body)
        # cell_parts = ['preamble', 'cell_name', 'cell_id', 'cell_content', ...]
        j = 1
        while j < len(cell_parts) - 2:
            cell_name = cell_parts[j]
            cell_content = cell_parts[j + 2].strip()
            cells[cell_name] = cell_content
            j += 3

        result[nb_name] = cells

    return result


# ---------------------------------------------------------------------------
# YAML comparison
# ---------------------------------------------------------------------------


def _load_yaml(path: Path) -> Optional[dict]:
    if yaml is None:
        return None
    if not path.exists():
        return None
    with open(path, encoding="utf-8") as f:
        return yaml.safe_load(f)


def _yaml_diff_table(old: Optional[dict], new: Optional[dict], keys: list[str]) -> str:
    """Produce a Markdown table comparing specific keys between two dicts."""
    rows = []
    for key in keys:
        old_val = old.get(key, "N/A") if old else "N/A"
        new_val = new.get(key, "N/A") if new else "N/A"
        changed = "Yes" if str(old_val) != str(new_val) else ""
        rows.append(f"| {key} | {old_val} | {new_val} | {changed} |")
    header = "| Parameter | Old | New | Changed |\n|-----------|-----|-----|---------|"
    return f"{header}\n" + "\n".join(rows)


def _deep_get(d: dict, *keys: str) -> Any:
    """Traverse nested dict with dotted keys."""
    for k in keys:
        for part in k.split("."):
            if isinstance(d, dict):
                d = d.get(part, {})
            else:
                return None
    return d if d != {} else None


# ---------------------------------------------------------------------------
# Metric extraction from manifest cells
# ---------------------------------------------------------------------------


def _extract_html_tables(cell_content: str) -> list:
    """Extract HTML tables from cell content using pandas.read_html."""
    if pd is None:
        return []
    # Find all <table>...</table> blocks
    tables = re.findall(r"<table[\s\S]*?</table>", cell_content, re.IGNORECASE)
    result = []
    for t in tables:
        try:
            dfs = pd.read_html(io.StringIO(t))
            result.extend(dfs)
        except Exception:
            pass
    return result


def _find_metric_tables(
    manifest: dict[str, dict[str, str]], nb_name: str = "08_baseline_experiments"
) -> list:
    """Find DataFrames that look like model-metric result tables in NB08."""
    cells = manifest.get(nb_name, {})
    result = []
    for cell_name, content in cells.items():
        for df in _extract_html_tables(content):
            cols = [c.lower() for c in df.columns]
            if any("auc" in c or "f1" in c for c in cols):
                result.append((cell_name, df))
    return result


# ---------------------------------------------------------------------------
# Cell-level diff
# ---------------------------------------------------------------------------


def _cell_diff(old_content: str, new_content: str, context: int = 3) -> str:
    """Produce a unified diff between two cell contents."""
    old_lines = old_content.split("\n")
    new_lines = new_content.split("\n")
    diff = difflib.unified_diff(old_lines, new_lines, lineterm="", n=context)
    return "\n".join(diff)


# ---------------------------------------------------------------------------
# Report generation
# ---------------------------------------------------------------------------


def _section_map_comparison(old_map: str, new_map: str) -> str:
    """Compare section maps from two manifests."""
    old_lines = [line.strip() for line in old_map.split("\n") if line.strip()]
    new_lines = [line.strip() for line in new_map.split("\n") if line.strip()]

    rows = []
    # Parse: <!-- stem (N cells, M with output) -->
    pat = re.compile(r"<!--\s*(\S+)\s+\((\d+)\s+cells,\s+(\d+)\s+with output\)\s*-->")
    old_parsed = {m.group(1): (int(m.group(2)), int(m.group(3))) for line in old_lines for m in [pat.match(line)] if m}
    new_parsed = {m.group(1): (int(m.group(2)), int(m.group(3))) for line in new_lines for m in [pat.match(line)] if m}

    all_nbs = sorted(set(old_parsed) | set(new_parsed))
    for nb in all_nbs:
        oc, oo = old_parsed.get(nb, (0, 0))
        nc, no = new_parsed.get(nb, (0, 0))
        cell_delta = nc - oc if oc and nc else "N/A"
        out_delta = no - oo if oc and nc else "N/A"
        rows.append(f"| {nb} | {oc} ({oo}) | {nc} ({no}) | {cell_delta} | {out_delta} |")

    header = "| Notebook | Old cells (outputs) | New cells (outputs) | Cell delta | Output delta |\n|----------|--------------------|--------------------|------------|--------------|"
    return f"{header}\n" + "\n".join(rows)


def _metrics_comparison(old_manifest: dict, new_manifest: dict) -> str:
    """Compare model metrics from NB08."""
    old_tables = _find_metric_tables(old_manifest)
    new_tables = _find_metric_tables(new_manifest)

    if not old_tables and not new_tables:
        return "_No metric tables found in either run._"

    parts = []
    for idx, (old_item, new_item) in enumerate(
        zip(old_tables or [(None, None)], new_tables or [(None, None)])
    ):
        old_name, old_df = old_item if old_item else (None, None)
        new_name, new_df = new_item if new_item else (None, None)

        if old_df is not None and new_df is not None:
            # Try to merge on model name column
            model_col = None
            for c in old_df.columns:
                if "model" in c.lower() or "name" in c.lower():
                    model_col = c
                    break

            if model_col:
                merged = old_df.merge(
                    new_df, on=model_col, suffixes=(" (old)", " (new)"), how="outer"
                )
                # Compute deltas for numeric columns
                for base_col in old_df.columns:
                    if base_col == model_col:
                        continue
                    old_c = f"{base_col} (old)"
                    new_c = f"{base_col} (new)"
                    if old_c in merged.columns and new_c in merged.columns:
                        try:
                            merged[f"{base_col} delta"] = (
                                merged[new_c].astype(float) - merged[old_c].astype(float)
                            ).round(4)
                        except (ValueError, TypeError):
                            pass

                parts.append(merged.to_markdown(index=False))

                # Flag regressions
                for base_col in old_df.columns:
                    delta_col = f"{base_col} delta"
                    if delta_col in merged.columns:
                        try:
                            min_delta = merged[delta_col].min()
                            if min_delta < -_REGRESSION_THRESHOLD:
                                parts.append(
                                    f"\n**REGRESSION**: {base_col} dropped by "
                                    f"{abs(min_delta):.4f} (threshold: {_REGRESSION_THRESHOLD})"
                                )
                            elif min_delta < -_DRIFT_THRESHOLD:
                                parts.append(
                                    f"\n**DRIFT**: {base_col} shifted by "
                                    f"{abs(min_delta):.4f} (threshold: {_DRIFT_THRESHOLD})"
                                )
                        except (ValueError, TypeError):
                            pass
            else:
                parts.append(f"**Old ({old_name}):**\n{old_df.to_markdown(index=False)}")
                parts.append(f"**New ({new_name}):**\n{new_df.to_markdown(index=False)}")
        elif old_df is not None:
            parts.append(f"**Old only ({old_name}):**\n{old_df.to_markdown(index=False)}")
        elif new_df is not None:
            parts.append(f"**New only ({new_name}):**\n{new_df.to_markdown(index=False)}")

    return "\n\n".join(parts) if parts else "_Could not parse metric tables._"


def _cell_by_cell_diff(old_manifest: dict, new_manifest: dict) -> str:
    """Produce cell-by-cell diffs for all notebooks."""
    parts = []
    all_nbs = sorted(
        set(old_manifest) | set(new_manifest) - {"__section_map__"}
    )

    for nb in all_nbs:
        if nb == "__section_map__":
            continue
        old_cells = old_manifest.get(nb, {})
        new_cells = new_manifest.get(nb, {})

        all_cell_names = list(dict.fromkeys(list(old_cells) + list(new_cells)))
        nb_diffs = []

        for cname in all_cell_names:
            old_c = old_cells.get(cname, "")
            new_c = new_cells.get(cname, "")
            if old_c == new_c:
                continue
            if not old_c:
                nb_diffs.append(f"#### {cname} (NEW)\n```\n{new_c[:500]}\n```")
            elif not new_c:
                nb_diffs.append(f"#### {cname} (REMOVED)\n```\n{old_c[:500]}\n```")
            else:
                diff = _cell_diff(old_c, new_c)
                if diff:
                    nb_diffs.append(f"#### {cname}\n```diff\n{diff}\n```")

        if nb_diffs:
            parts.append(f"### {nb}\n\n" + "\n\n".join(nb_diffs))

    return "\n\n".join(parts) if parts else "_No cell-level differences detected._"


def _yaml_section(
    title: str,
    old_dir: Path,
    new_dir: Path,
    filename: str,
    keys: list[str],
    sub_path: str = "",
) -> str:
    """Compare a specific YAML file between two run directories."""
    old_path = old_dir / sub_path / filename if sub_path else old_dir / filename
    new_path = new_dir / sub_path / filename if sub_path else new_dir / filename
    old_data = _load_yaml(old_path)
    new_data = _load_yaml(new_path)

    if old_data is None and new_data is None:
        return f"## {title}\n\n_File not found in either run: {filename}_"

    parts = [f"## {title}"]
    parts.append(_yaml_diff_table(old_data, new_data, keys))

    # Special handling for grid_dates count
    if filename == "snapshot_grid.yaml":
        old_dates = old_data.get("grid_dates", []) if old_data else []
        new_dates = new_data.get("grid_dates", []) if new_data else []
        parts.append(
            f"\nGrid dates: {len(old_dates)} (old) vs {len(new_dates)} (new)"
        )
        if old_dates and new_dates:
            if old_dates[0] != new_dates[0] or old_dates[-1] != new_dates[-1]:
                parts.append(
                    f"Range: [{old_dates[0]} .. {old_dates[-1]}] → "
                    f"[{new_dates[0]} .. {new_dates[-1]}]"
                )

    return "\n".join(parts)


def _feature_profile_comparison(old_dir: Path, new_dir: Path) -> str:
    """Compare exploration_feature_profile.yaml between runs."""
    old_fp = _load_yaml(old_dir / "merged" / "exploration_feature_profile.yaml")
    new_fp = _load_yaml(new_dir / "merged" / "exploration_feature_profile.yaml")

    if old_fp is None and new_fp is None:
        return "## Feature Profile\n\n_Not found in either run._"

    parts = ["## Feature Profile"]

    old_features = set(old_fp.keys()) if old_fp else set()
    new_features = set(new_fp.keys()) if new_fp else set()

    added = new_features - old_features
    removed = old_features - new_features
    common = old_features & new_features

    parts.append(f"- Old features: {len(old_features)}")
    parts.append(f"- New features: {len(new_features)}")
    parts.append(f"- Added: {len(added)}")
    parts.append(f"- Removed: {len(removed)}")
    parts.append(f"- Common: {len(common)}")

    if added:
        parts.append(f"\n**Added features:** {', '.join(sorted(added)[:20])}")
        if len(added) > 20:
            parts.append(f"  ...and {len(added) - 20} more")
    if removed:
        parts.append(f"\n**Removed features:** {', '.join(sorted(removed)[:20])}")
        if len(removed) > 20:
            parts.append(f"  ...and {len(removed) - 20} more")

    return "\n".join(parts)


def _timing_comparison(old_log: Path, new_log: Path) -> str:
    """Compare per-notebook timing from run_exploration.log files."""
    parts = ["## Per-Notebook Timing"]

    def _parse_log(path: Path) -> dict[str, str]:
        if not path.exists():
            return {}
        timings = {}
        for line in path.read_text(encoding="utf-8").split("\n"):
            m = re.match(r"\s*(\S+)\s+.*?OK\s+\((\d+\.?\d*)s\)", line)
            if m:
                timings[m.group(1)] = m.group(2)
            else:
                m = re.match(r"\s*(\S+)\s+.*?(FAILED|SKIPPED)", line)
                if m:
                    timings[m.group(1)] = m.group(2)
        return timings

    old_t = _parse_log(old_log)
    new_t = _parse_log(new_log)

    if not old_t and not new_t:
        return parts[0] + "\n\n_No timing logs found._"

    all_nbs = list(dict.fromkeys(list(old_t) + list(new_t)))
    rows = []
    for nb in all_nbs:
        ov = old_t.get(nb, "N/A")
        nv = new_t.get(nb, "N/A")
        try:
            delta = f"{float(nv) - float(ov):+.1f}s"
        except (ValueError, TypeError):
            delta = "N/A"
        rows.append(f"| {nb} | {ov} | {nv} | {delta} |")

    header = "| Notebook | Old | New | Delta |\n|----------|-----|-----|-------|"
    parts.append(f"{header}\n" + "\n".join(rows))
    return "\n".join(parts)


def _read_run_metadata(run_dir: Optional[Path]) -> dict[str, str]:
    """Read git_sha.txt, version.txt, and commit_date.txt from a run directory."""
    meta: dict[str, str] = {"sha": "unknown", "version": "unknown", "date": "unknown"}
    if run_dir is None:
        return meta
    for key, filename in [("sha", "git_sha.txt"), ("version", "version.txt"), ("date", "commit_date.txt")]:
        p = run_dir / filename
        if p.exists():
            meta[key] = p.read_text(encoding="utf-8").strip()
    # Truncate SHA to 7 chars for display
    if meta["sha"] != "unknown":
        meta["sha"] = meta["sha"][:7]
    return meta


def generate_drift_report(
    old_manifest_path: Path,
    new_manifest_path: Path,
    old_run_dir: Optional[Path],
    new_run_dir: Optional[Path],
) -> str:
    """Generate the full drift report."""
    old_text = old_manifest_path.read_text(encoding="utf-8")
    new_text = new_manifest_path.read_text(encoding="utf-8")

    old_m = _parse_manifest(old_text)
    new_m = _parse_manifest(new_text)

    old_meta = _read_run_metadata(old_run_dir)
    new_meta = _read_run_metadata(new_run_dir)

    old_label = f"{old_meta['sha']}  v{old_meta['version']}  ({old_meta['date']})"
    new_label = f"{new_meta['sha']}  v{new_meta['version']}  ({new_meta['date']})"

    sections = []

    # Title
    sections.append(f"# Drift Report: {old_meta['sha']} → {new_meta['sha']}")

    # Summary
    old_nb_count = len([k for k in old_m if k != "__section_map__"])
    new_nb_count = len([k for k in new_m if k != "__section_map__"])
    sections.append(
        f"## Summary\n"
        f"- Old: {old_label}\n"
        f"- New: {new_label}\n"
        f"- Notebooks captured: {old_nb_count} (old), {new_nb_count} (new)"
    )

    # Section map changes
    sections.append("## Section Map Changes\n")
    sections.append(
        _section_map_comparison(
            old_m.get("__section_map__", {}).get("_raw", ""),
            new_m.get("__section_map__", {}).get("_raw", ""),
        )
    )

    # Model metrics (NB08) — the critical section
    sections.append("## Model Metrics (NB08) — Critical\n")
    sections.append(_metrics_comparison(old_m, new_m))

    # YAML comparisons
    if old_run_dir and new_run_dir:
        sections.append(
            _yaml_section(
                "Snapshot Grid",
                old_run_dir,
                new_run_dir,
                "snapshot_grid.yaml",
                [
                    "grid_start",
                    "grid_end",
                    "cadence_interval",
                    "observation_window_days",
                    "purge_gap_days",
                    "label_window_days",
                    "max_grid_dates",
                ],
            )
        )

        sections.append(
            _yaml_section(
                "Project Context / Intent",
                old_run_dir,
                new_run_dir,
                "project_context.yaml",
                [
                    "prediction_horizons",
                    "recent_window_days",
                    "observation_window_days",
                    "purge_gap_days",
                    "label_window_days",
                    "split_strategy",
                    "cadence_interval",
                    "temporal_posture",
                ],
            )
        )

        sections.append(_feature_profile_comparison(old_run_dir, new_run_dir))

        # Timing
        old_log = old_run_dir.parent.parent / "exploration_notebooks" / "run_exploration.log"
        new_log = new_run_dir.parent.parent / "exploration_notebooks" / "run_exploration.log"
        # Also check in the parent of run_dir
        if not old_log.exists():
            old_log = old_run_dir.parent.parent / "run_exploration.log"
        if not new_log.exists():
            new_log = new_run_dir.parent.parent / "run_exploration.log"
        sections.append(_timing_comparison(old_log, new_log))

    # Cell-by-cell diff
    sections.append("## Cell-by-Cell Diff\n")
    sections.append(_cell_by_cell_diff(old_m, new_m))

    return "\n\n".join(sections) + "\n"


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


def main() -> None:
    parser = argparse.ArgumentParser(description="Generate exploration drift report.")
    parser.add_argument("--old-manifest", type=Path, required=True)
    parser.add_argument("--new-manifest", type=Path, required=True)
    parser.add_argument("--old-run-dir", type=Path, default=None)
    parser.add_argument("--new-run-dir", type=Path, default=None)
    parser.add_argument("--output", type=Path, default=None, help="Output file (default: stdout)")
    args = parser.parse_args()

    report = generate_drift_report(
        args.old_manifest, args.new_manifest, args.old_run_dir, args.new_run_dir
    )

    if args.output:
        args.output.parent.mkdir(parents=True, exist_ok=True)
        args.output.write_text(report, encoding="utf-8")
        print(f"Drift report written to {args.output}", file=sys.stderr)
    else:
        print(report)


if __name__ == "__main__":
    main()
