"""
REASONING LAB
=============
Autonomous training ground for Harvester's reasoning engine.
Feeds hard architectural problems through reason() → run_engine(),
banking reasoning plans, patterns, and failure lessons.

The code follows the reasoning. Not the other way around.

Usage:
  python reasoning_lab.py                 # run all tasks
  python reasoning_lab.py --task 5        # run only task #5
  python reasoning_lab.py --list          # list all tasks
  python reasoning_lab.py --reason-only   # reason without coding (thinking drills)
  python reasoning_lab.py --difficulty 4  # only tasks with difficulty >= 4
  python reasoning_lab.py --category ai   # only tasks in a specific category
"""

import os
import sys
import json
import time
import argparse

sys.path.insert(0, os.path.dirname(__file__))

from growth_engine import run_growth_engine, init_db, show_stats
from reflection_engine import reason, set_language, set_task_mode, ReasoningPlan

# ── Config ────────────────────────────────────────────────────────────────────

DEFAULT_TASK_BANK = os.path.join(os.path.dirname(__file__), "reasoning_tasks.json")
RESULTS_DIR = os.path.join(os.path.dirname(__file__), "lab_results")

# ── Load tasks ────────────────────────────────────────────────────────────────

def load_tasks(
    task_id: int = None,
    min_difficulty: int = None,
    category: str = None,
    bank_path: str = None
) -> list[dict]:
    path = bank_path or DEFAULT_TASK_BANK
    with open(path, "r", encoding="utf-8") as f:
        bank = json.load(f)

    tasks = bank["tasks"] if isinstance(bank, dict) else bank

    # Normalize field names across bank formats
    for t in tasks:
        if "prompt" not in t and "description" in t:
            t["prompt"] = t["description"]
        if "lang" not in t and "language" in t:
            t["lang"] = t["language"]

    if task_id is not None:
        tasks = [t for t in tasks if t["id"] == task_id]
    if min_difficulty is not None:
        tasks = [t for t in tasks if t["difficulty"] >= min_difficulty]
    if category is not None:
        tasks = [t for t in tasks if t["category"] == category.lower()]

    return tasks


# ── Reason-only mode ──────────────────────────────────────────────────────────

def run_reason_only(tasks: list[dict]) -> list[dict]:
    """
    Pure thinking drills. No code generation.
    Exercises reason() and banks the plans for review.
    """
    results = []
    for i, task in enumerate(tasks, 1):
        print(f"\n{'='*60}")
        print(f"  REASONING DRILL {i}/{len(tasks)}")
        print(f"  [{task['difficulty']}★] {task['name']}")
        print(f"  Category: {task['category']}")
        print(f"{'='*60}")

        set_language(task["lang"])
        t0 = time.time()
        plan = reason(task["prompt"])
        elapsed = time.time() - t0

        result = {
            "task_id":       task["id"],
            "name":          task["name"],
            "difficulty":    task["difficulty"],
            "category":      task["category"],
            "approach":      plan.approach,
            "sub_problems":  plan.sub_problems,
            "steps":         plan.steps,
            "edge_cases":    plan.edge_cases,
            "confidence":    plan.confidence,
            "time_secs":     round(elapsed, 1)
        }
        results.append(result)

        print(f"\n  ⏱  Reasoning took {elapsed:.1f}s")
        print(f"  📊 Confidence: {plan.confidence:.0%}")
        print(f"  🧩 Sub-problems: {len(plan.sub_problems)}")
        print(f"  📝 Steps: {len(plan.steps)}")
        print(f"  ⚠️  Edge cases: {len(plan.edge_cases)}")

    return results


# ── Full lab mode ─────────────────────────────────────────────────────────────

def run_full_lab(tasks: list[dict]) -> list[dict]:
    """
    Full pipeline: reason → generate → reflect → score → learn.
    Every run makes the engine permanently smarter.
    """
    results = []
    for i, task in enumerate(tasks, 1):
        print(f"\n{'='*60}")
        print(f"  REASONING LAB — TASK {i}/{len(tasks)}")
        print(f"  [{task['difficulty']}★] {task['name']}")
        print(f"  Category: {task['category']}")
        print(f"{'='*60}")

        t0 = time.time()
        attempt = run_growth_engine(task["prompt"], lang=task["lang"])
        elapsed = time.time() - t0

        result = {
            "task_id":      task["id"],
            "name":         task["name"],
            "difficulty":   task["difficulty"],
            "category":     task["category"],
            "score":        attempt.score,
            "passed":       attempt.passed,
            "iterations":   attempt.iteration,
            "time_secs":    round(elapsed, 1),
            "reflections":  attempt.reflections[:5],
            "errors":       attempt.errors[:3]
        }
        results.append(result)

        status = "PASS" if attempt.passed else "FAIL"
        print(f"\n  {'='*40}")
        print(f"  {status} — Score: {attempt.score:.1f}% in {attempt.iteration} iterations ({elapsed:.0f}s)")
        print(f"  {'='*40}")

    return results


# ── Report ────────────────────────────────────────────────────────────────────

def save_report(results: list[dict], mode: str):
    os.makedirs(RESULTS_DIR, exist_ok=True)
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    filename = f"lab_{mode}_{timestamp}.json"
    filepath = os.path.join(RESULTS_DIR, filename)

    report = {
        "mode":       mode,
        "timestamp":  timestamp,
        "task_count": len(results),
        "results":    results
    }

    # Add summary stats
    if mode == "full":
        scores = [r["score"] for r in results]
        passed = sum(1 for r in results if r["passed"])
        report["summary"] = {
            "passed":     passed,
            "failed":     len(results) - passed,
            "avg_score":  round(sum(scores) / len(scores), 1) if scores else 0,
            "best_score": max(scores) if scores else 0,
            "worst_score": min(scores) if scores else 0,
            "total_time": round(sum(r["time_secs"] for r in results), 1)
        }
    elif mode == "reason":
        confs = [r["confidence"] for r in results]
        report["summary"] = {
            "avg_confidence": round(sum(confs) / len(confs), 2) if confs else 0,
            "high_confidence": sum(1 for c in confs if c >= 0.7),
            "low_confidence":  sum(1 for c in confs if c < 0.5),
            "total_time": round(sum(r["time_secs"] for r in results), 1)
        }

    with open(filepath, "w", encoding="utf-8") as f:
        json.dump(report, f, indent=2)

    print(f"\n  📄 Report saved: {filepath}")
    return report


def print_summary(report: dict):
    s = report.get("summary", {})
    mode = report["mode"]
    print(f"\n{'='*60}")
    print(f"  REASONING LAB — SUMMARY ({mode.upper()} MODE)")
    print(f"{'='*60}")
    print(f"  Tasks run: {report['task_count']}")

    if mode == "full":
        print(f"  Passed:    {s.get('passed', 0)}")
        print(f"  Failed:    {s.get('failed', 0)}")
        print(f"  Avg score: {s.get('avg_score', 0)}%")
        print(f"  Best:      {s.get('best_score', 0)}%")
        print(f"  Worst:     {s.get('worst_score', 0)}%")
    elif mode == "reason":
        print(f"  Avg confidence:  {s.get('avg_confidence', 0):.0%}")
        print(f"  High confidence: {s.get('high_confidence', 0)} tasks")
        print(f"  Low confidence:  {s.get('low_confidence', 0)} tasks")

    print(f"  Total time: {s.get('total_time', 0)}s")
    print(f"{'='*60}")


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

def main():
    parser = argparse.ArgumentParser(description="Reasoning Lab — train Harvester's thinking")
    parser.add_argument("--task", type=int, help="Run a specific task by ID")
    parser.add_argument("--list", action="store_true", help="List all tasks")
    parser.add_argument("--reason-only", action="store_true", help="Reason without coding")
    parser.add_argument("--difficulty", type=int, help="Only tasks with difficulty >= N")
    parser.add_argument("--category", type=str, help="Only tasks in this category")
    parser.add_argument("--bank", type=str, help="Path to task bank JSON (default: reasoning_tasks.json)")
    args = parser.parse_args()

    # List mode
    if args.list:
        tasks = load_tasks(bank_path=args.bank)
        print(f"\n  REASONING LAB — {len(tasks)} TASKS")
        print(f"  {'='*50}")
        for t in tasks:
            print(f"  [{t['id']:2d}] [{t['difficulty']}★] {t['name']:<35} ({t['category']})")
        categories = sorted(set(t["category"] for t in tasks))
        print(f"\n  Categories: {', '.join(categories)}")
        return

    # Init database
    init_db()

    # Load filtered tasks
    tasks = load_tasks(
        task_id=args.task,
        min_difficulty=args.difficulty,
        category=args.category,
        bank_path=args.bank
    )

    if not tasks:
        print("  No tasks matched your filters.")
        return

    print(f"\n  REASONING LAB — {len(tasks)} task(s) queued")

    # Auto-detect task mode from bank name: OSINT/Bellingcat → DeepSeek R1, else → Qwen
    bank_name = os.path.basename(args.bank or "").lower()
    if any(kw in bank_name for kw in ("osint", "bellingcat", "investigation", "missing")):
        set_task_mode("reasoning")
    else:
        set_task_mode("code")

    # Run
    if args.reason_only:
        results = run_reason_only(tasks)
        report = save_report(results, "reason")
    else:
        results = run_full_lab(tasks)
        report = save_report(results, "full")

    print_summary(report)

    # Show growth stats
    print()
    show_stats()


if __name__ == "__main__":
    main()