#!/usr/bin/env python3
from __future__ import annotations
"""
Phase 15: Auto-Lab Orchestrator (Safe Mode Default)
---------------------------------------------------
Coordinates multi-phase substrate analysis without making
live BloomOS calls unless --live is provided.

Pipeline:
    1. Load pheromones
    2. Multi-task co-evolver (best configs per task)
    3. (optional) Auto-experimenter (live experiments with --live)
    4. Evolution engine (mutations & drift detection)
    5. Capsule evolver (semantic archetypes)
    6. Produce unified report

Usage (safe mode - no BloomOS calls):
    PYTHONPATH=. python3 -m helix_cdc.bloomos.auto_lab \\
        /tmp/pheromones.jsonl \\
        --out /tmp/auto_lab_report.json \\
        --pretty

Usage (live mode - runs experiments):
    PYTHONPATH=. python3 -m helix_cdc.bloomos.auto_lab \\
        /tmp/pheromones.jsonl \\
        --capsules /tmp/capsules.jsonl \\
        --out /tmp/auto_lab_report.json \\
        --live \\
        --goal "optimize math performance" \\
        --pretty
"""

import argparse
import json
from dataclasses import asdict
from pathlib import Path
from typing import Any, Dict, List, Optional

# Phase imports
from .multitask_coevolver import analyze_multitask, MultiTaskPlan, TaskConfig
from .evolution_engine import EvolutionEngine, EvolutionResult
from .capsule_evolver import analyze_capsules, CapsuleEvolutionSummary


# =========================================================
# Helper: Extract recommendations from MultiTaskPlan
# =========================================================

def _extract_recommendations(plan: MultiTaskPlan) -> Dict[str, Any]:
    """
    Extract per-task recommendations from MultiTaskPlan.
    Returns dict keyed by task_type with best config for each.
    """
    recommendations: Dict[str, Any] = {}

    for cfg in plan.task_configs:
        recommendations[cfg.task_type] = {
            "gate": cfg.gate,
            "helix_program": cfg.helix_program,
            "regrow_depth": cfg.regrow_depth,
            "lattice_D": cfg.lattice_D,
            "avg_fitness": cfg.avg_fitness,
            "avg_Se": cfg.avg_Se,
            "num_samples": cfg.num_samples,
            "notes": cfg.notes,
        }

    return recommendations


def _extract_best_configs(plan: MultiTaskPlan) -> Dict[str, List[Any]]:
    """
    Extract unique best gates, programs, depths, lattice_Ds across all tasks.
    """
    gates = set()
    programs = set()
    depths = set()
    lattice_Ds = set()

    for cfg in plan.task_configs:
        if cfg.gate:
            gates.add(cfg.gate)
        if cfg.helix_program:
            programs.add(cfg.helix_program)
        if cfg.regrow_depth:
            depths.add(cfg.regrow_depth)
        if cfg.lattice_D:
            lattice_Ds.add(cfg.lattice_D)

    return {
        "gates": sorted(gates),
        "programs": sorted(programs),
        "depths": sorted(depths),
        "lattice_Ds": sorted(lattice_Ds),
    }


# =========================================================
# Auto-Lab Result Bundle
# =========================================================

def _bundle_results(
    source_path: str,
    multitask_plan: MultiTaskPlan,
    evolution_results: List[EvolutionResult],
    capsule_summary: Optional[CapsuleEvolutionSummary],
    experiments: Optional[Dict[str, Any]],
    live_mode: bool,
) -> Dict[str, Any]:
    """Bundle all results into a single JSON-serializable dict."""

    # Convert MultiTaskPlan
    multitask_dict = {
        "path": multitask_plan.path,
        "total_entries": multitask_plan.total_entries,
        "task_configs": [asdict(cfg) for cfg in multitask_plan.task_configs],
    }

    # Convert EvolutionResults
    evolution_dict = [asdict(r) for r in evolution_results]

    # Convert CapsuleEvolutionSummary
    if capsule_summary:
        capsule_dict = {
            "source_path": capsule_summary.source_path,
            "total_records": capsule_summary.total_records,
            "task_types": capsule_summary.task_types,
            "archetypes": [asdict(a) for a in capsule_summary.archetypes],
        }
    else:
        capsule_dict = {"note": "No capsule file provided or no capsules found"}

    # Extract recommendations
    recommendations = _extract_recommendations(multitask_plan)
    best_configs = _extract_best_configs(multitask_plan)

    return {
        "source_pheromones": source_path,
        "mode": "live" if live_mode else "safe",
        "recommendations": recommendations,
        "best_configs": best_configs,
        "multitask_analysis": multitask_dict,
        "evolution_results": evolution_dict,
        "capsule_archetypes": capsule_dict,
        "experiment_results": experiments if experiments else {"mode": "dry", "details": None},
    }


# =========================================================
# Auto-Lab Core
# =========================================================

def run_auto_lab(
    pheromone_path: str,
    capsule_path: Optional[str] = None,
    out_path: Optional[str] = None,
    min_samples: int = 1,
    live: bool = False,
    goal: str = "auto-lab evaluation",
    limit_experiments: Optional[int] = None,
    window: int = 50,
    prev_window: int = 25,
    mutation_rate: float = 0.12,
    drift_threshold: float = -0.05,
) -> Dict[str, Any]:
    """
    Run the full Auto-Lab pipeline.

    Args:
        pheromone_path: Path to pheromone JSONL file
        capsule_path: Optional path to capsule JSONL file
        out_path: Optional output JSON path
        min_samples: Minimum samples for coevolver buckets
        live: If True, run live experiments (requires BloomOS runtime)
        goal: Goal string for experiments
        limit_experiments: Max experiments to run in live mode
        window: Evolution engine current window size
        prev_window: Evolution engine previous window size
        mutation_rate: Evolution engine mutation rate
        drift_threshold: Drift threshold to trigger mutation

    Returns:
        Bundled results dict
    """
    pheromone_path = str(Path(pheromone_path).resolve())

    # -------------------------------
    # 1. Multitask co-evolution scan
    # -------------------------------
    multitask_plan = analyze_multitask(
        pheromone_path,
        min_samples=min_samples,
    )

    # -------------------------------
    # 2. Experiments (optional live)
    # -------------------------------
    experiments_dict: Optional[Dict[str, Any]] = None

    if live:
        try:
            from .auto_experimenter import AutoExperimenter, ExperimentConfig

            best = _extract_best_configs(multitask_plan)

            experimenter = AutoExperimenter(
                pheromone_path=pheromone_path,
                base_url="http://localhost:7070",
            )

            results, summary = experimenter.run_sweep(
                goal=goal,
                task_type="experiment",
                gates=best["gates"] if best["gates"] else None,
                programs=best["programs"] if best["programs"] else None,
                regrow_depths=best["depths"] if best["depths"] else None,
                lattice_Ds=best["lattice_Ds"] if best["lattice_Ds"] else None,
                limit=limit_experiments,
            )

            experiments_dict = {
                "mode": "live",
                "goal": goal,
                "total_experiments": len(results),
                "summary": asdict(summary) if summary else None,
                "results": [asdict(r) for r in results],
            }
        except Exception as e:
            experiments_dict = {
                "mode": "live",
                "error": str(e),
                "note": "Live experiments failed - is BloomOS runtime running?",
            }

    # -------------------------------
    # 3. Evolution engine (mutations)
    # -------------------------------
    engine = EvolutionEngine(
        pheromone_path=pheromone_path,
        window=window,
        prev_window=prev_window,
        mutation_rate=mutation_rate,
        min_samples=min_samples,
        drift_threshold=drift_threshold,
    )
    evolution_results = engine.evolve_all_tasks()

    # -------------------------------
    # 4. Capsule analysis
    # -------------------------------
    capsule_summary: Optional[CapsuleEvolutionSummary] = None

    if capsule_path and Path(capsule_path).exists():
        capsule_summary = analyze_capsules(
            path=capsule_path,
            min_cluster_size=1,
        )

    # -------------------------------
    # 5. Bundle results
    # -------------------------------
    bundle = _bundle_results(
        source_path=pheromone_path,
        multitask_plan=multitask_plan,
        evolution_results=evolution_results,
        capsule_summary=capsule_summary,
        experiments=experiments_dict,
        live_mode=live,
    )

    # -------------------------------
    # 6. Write output if requested
    # -------------------------------
    if out_path:
        out_file = Path(out_path)
        out_file.parent.mkdir(parents=True, exist_ok=True)
        out_file.write_text(json.dumps(bundle, indent=2))

    return bundle


# =========================================================
# Pretty Print
# =========================================================

def _print_pretty(bundle: Dict[str, Any]) -> None:
    """Print human-readable summary of Auto-Lab results."""
    print("=" * 70)
    print("AUTO-LAB ORCHESTRATOR - Phase 15")
    print("=" * 70)
    print(f"Source: {bundle['source_pheromones']}")
    print(f"Mode:   {bundle['mode']}")
    print()

    # Recommendations
    print("-" * 70)
    print("RECOMMENDATIONS (per task_type)")
    print("-" * 70)
    recs = bundle.get("recommendations", {})
    if not recs:
        print("  (no recommendations)")
    else:
        print(f"{'Task':<12} {'Gate':<12} {'Program':<10} {'Depth':>6} {'D':>4} {'Fitness':>10}")
        print("-" * 70)
        for task, cfg in sorted(recs.items()):
            gate = cfg.get("gate") or "-"
            prog = cfg.get("helix_program") or "-"
            depth = cfg.get("regrow_depth") or "-"
            D = cfg.get("lattice_D") or "-"
            fit = cfg.get("avg_fitness")
            fit_str = f"{fit:.2f}" if fit else "-"
            print(f"{task:<12} {gate:<12} {prog:<10} {str(depth):>6} {str(D):>4} {fit_str:>10}")
    print()

    # Evolution results
    print("-" * 70)
    print("EVOLUTION ENGINE RESULTS")
    print("-" * 70)
    evo = bundle.get("evolution_results", [])
    if not evo:
        print("  (no evolution results)")
    else:
        for r in evo:
            task = r.get("task_type", "?")
            drift = r.get("drift")
            drift_str = f"{drift:.4f}" if drift else "N/A"
            mutated = r.get("mutation_applied", False)
            rec_prog = r.get("recommended_program", "-")
            rec_gate = r.get("recommended_gate", "-")
            print(f"  {task}: drift={drift_str}, mutated={mutated}, rec_gate={rec_gate}, rec_prog={rec_prog}")
    print()

    # Capsule archetypes summary
    print("-" * 70)
    print("CAPSULE ARCHETYPES")
    print("-" * 70)
    caps = bundle.get("capsule_archetypes", {})
    archetypes = caps.get("archetypes", [])
    if not archetypes:
        note = caps.get("note", "No archetypes")
        print(f"  {note}")
    else:
        print(f"  Total archetypes: {len(archetypes)}")
        print(f"  Task types: {', '.join(caps.get('task_types', []))}")
        # Show top 5 by fitness
        sorted_arch = sorted(archetypes, key=lambda a: -(a.get("avg_fitness") or 0))[:5]
        print(f"\n  Top 5 by fitness:")
        for a in sorted_arch:
            aid = a.get("id", "?")
            fit = a.get("avg_fitness")
            fit_str = f"{fit:.2f}" if fit else "-"
            kws = ", ".join(a.get("keywords", [])[:5])
            print(f"    {aid}: fitness={fit_str}, keywords=[{kws}]")
    print()

    # Experiments
    print("-" * 70)
    print("EXPERIMENTS")
    print("-" * 70)
    exp = bundle.get("experiment_results", {})
    mode = exp.get("mode", "dry")
    if mode == "dry":
        print("  Mode: dry (no live experiments)")
    elif "error" in exp:
        print(f"  Mode: live (FAILED)")
        print(f"  Error: {exp.get('error')}")
    else:
        total = exp.get("total_experiments", 0)
        print(f"  Mode: live")
        print(f"  Total experiments: {total}")
        summary = exp.get("summary")
        if summary:
            print(f"  Best fitness: {summary.get('best_fitness', 'N/A')}")
            print(f"  Best program: {summary.get('best_program', 'N/A')}")
    print()


# =========================================================
# CLI
# =========================================================

def _cli() -> None:
    parser = argparse.ArgumentParser(
        prog="auto_lab",
        description="Phase 15: Auto-Lab Orchestrator (Safe Mode Default)",
    )

    parser.add_argument(
        "pheromone_path",
        help="Pheromone JSONL file path",
    )
    parser.add_argument(
        "--capsules",
        dest="capsule_path",
        help="Capsule JSONL file path (optional)",
    )
    parser.add_argument(
        "--out",
        dest="out_path",
        help="Output JSON file path",
    )
    parser.add_argument(
        "--min-samples",
        type=int,
        default=1,
        help="Minimum samples per bucket (default: 1)",
    )
    parser.add_argument(
        "--live",
        action="store_true",
        help="Enable live experiments (requires BloomOS runtime)",
    )
    parser.add_argument(
        "--goal",
        default="auto-lab evaluation",
        help="Goal string for experiments",
    )
    parser.add_argument(
        "--limit-experiments",
        type=int,
        help="Max experiments to run in live mode",
    )
    parser.add_argument(
        "--window",
        type=int,
        default=50,
        help="Evolution engine current window (default: 50)",
    )
    parser.add_argument(
        "--prev-window",
        type=int,
        default=25,
        help="Evolution engine previous window (default: 25)",
    )
    parser.add_argument(
        "--mutation-rate",
        type=float,
        default=0.12,
        help="Mutation rate (default: 0.12)",
    )
    parser.add_argument(
        "--drift-threshold",
        type=float,
        default=-0.05,
        help="Drift threshold (default: -0.05)",
    )
    parser.add_argument(
        "--json",
        action="store_true",
        help="Output JSON to stdout",
    )
    parser.add_argument(
        "--pretty",
        action="store_true",
        help="Pretty-print human-readable summary",
    )

    args = parser.parse_args()

    bundle = run_auto_lab(
        pheromone_path=args.pheromone_path,
        capsule_path=args.capsule_path,
        out_path=args.out_path,
        min_samples=args.min_samples,
        live=args.live,
        goal=args.goal,
        limit_experiments=args.limit_experiments,
        window=args.window,
        prev_window=args.prev_window,
        mutation_rate=args.mutation_rate,
        drift_threshold=args.drift_threshold,
    )

    if args.json:
        print(json.dumps(bundle, indent=2))

    if args.pretty or not (args.json or args.out_path):
        _print_pretty(bundle)


if __name__ == "__main__":
    _cli()
