"""
s-LLM-project: Complete Experiment Runner

Runs all phases of the graph-enhanced reasoning experiment:
  Phase 1: Baselines — IO, CoT, CoT-SC, ToT
  Phase 2: AGoT — Adaptive Graph of Thoughts

Usage:
  python src/run_all.py --phase all --samples 5
  python src/run_all.py --phase baseline --task gsm8k --samples 10
  python src/run_all.py --phase agot --task math --samples 3
"""

import json
import os
import sys
import time
from pathlib import Path

# Add project root to path
_project_root = Path(__file__).resolve().parent.parent
if str(_project_root) not in sys.path:
    sys.path.insert(0, str(_project_root))

from src.samples import get_dataset, get_all_small
from src.baseline.benchmark import BaselineRunner, extract_answer, is_correct
from src.agot.run import evaluate_agot


def run_phase_baseline(
    model: str = "qwen2.5:32b",
    tasks: list = None,
    samples: int = -1,
    methods: list = None,
    results_dir: str = "data/results",
) -> dict:
    """Phase 1: Run baseline evaluations."""
    if tasks is None:
        tasks = ["gsm8k", "math", "aime"]
    if methods is None:
        methods = ["IO", "CoT"]

    runner = BaselineRunner(model=model, results_dir=results_dir)
    all_results = {}

    for task in tasks:
        dataset = get_dataset(task, samples)
        print(f"\n{'#'*70}")
        print(f"# Phase 1 — {model} on {task} ({len(dataset)} samples)")
        print(f"{'#'*70}")
        task_results = runner.run_all(dataset, task_name=task, methods=methods)
        all_results[task] = task_results

    return all_results


def run_phase_agot(
    model: str = "qwen2.5:32b",
    tasks: list = None,
    samples: int = -1,
    max_depth: int = 2,
    max_layers: int = 2,
    results_dir: str = "data/results",
) -> dict:
    """Phase 2: Run AGoT evaluation."""
    if tasks is None:
        tasks = ["gsm8k", "math"]

    all_results = {}

    for task in tasks:
        dataset = get_dataset(task, samples)
        print(f"\n{'#'*70}")
        print(f"# Phase 2 — AGoT ({model}) on {task} ({len(dataset)} samples)")
        print(f"{'#'*70}")
        result = evaluate_agot(
            dataset, task, model,
            max_depth=max_depth, max_layers=max_layers,
            results_dir=results_dir,
        )
        all_results[task] = result

    return all_results


def print_master_summary(all_results: dict, phase_name: str):
    """Print a formatted summary table."""
    print(f"\n\n{'='*70}")
    print(f"  Master Summary — {phase_name}")
    print(f"{'='*70}")
    print(f"{'Task':<15} {'Method':<22} {'Accuracy':<12} {'Time(s)':<10}")
    print(f"{'-'*60}")

    for task, results in all_results.items():
        if isinstance(results, dict):
            # Could be nested (task -> method -> summary)
            first = True
            for method, res in results.items():
                acc = f"{res.get('accuracy', 0):.2%}" if isinstance(res, dict) else "N/A"
                t = str(res.get("avg_time_s", "N/A")) if isinstance(res, dict) else "N/A"
                task_label = task if first else ""
                print(f"{task_label:<15} {method:<22} {acc:<12} {t:<10}")
                first = False


# ── CLI ─────────────────────────────────────────────────────────────

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="s-LLM-project experiment runner")
    parser.add_argument("--phase", choices=["baseline", "agot", "all"], default="all")
    parser.add_argument("--model", default="qwen2.5:32b",
                        help="Model name (qwen2.5:32b, qwq:latest, etc.)")
    parser.add_argument("--task", choices=["gsm8k", "math", "aime", "c_gsm8k", "cmath", "c_aime", "all"], default="all")
    parser.add_argument("--samples", type=int, default=-1,
                        help="Number of samples (-1 = all available)")
    parser.add_argument("--max-depth", type=int, default=2, help="AGoT max depth")
    parser.add_argument("--max-layers", type=int, default=2, help="AGoT max layers")
    parser.add_argument("--all-models", action="store_true",
                        help="Run across all available models")
    args = parser.parse_args()

    models = [args.model]
    if args.all_models:
        models = ["qwen2.5:32b", "qwq:latest", "qwen2.5:14b"]

    tasks = ["gsm8k", "math", "aime"] if args.task == "all" else [args.task]

    for model in models:
        print(f"\n{'='*70}")
        print(f"  Model: {model}")
        print(f"{'='*70}")

        if args.phase in ("baseline", "all"):
            bl_results = run_phase_baseline(model, tasks, args.samples)
            print_master_summary(bl_results, f"Phase 1 — Baseline — {model}")

        if args.phase in ("agot", "all"):
            # AGoT is expensive, use fewer samples
            agot_samples = min(args.samples, 5) if args.samples > 0 else 5
            agot_tasks = [t for t in tasks if t != "aime"]  # skip hardest for speed
            agot_results = run_phase_agot(
                model, agot_tasks, agot_samples,
                max_depth=args.max_depth, max_layers=args.max_layers,
            )
            print_master_summary(agot_results, f"Phase 2 — AGoT — {model}")

    print("\n✅ All experiments complete!")
