import os
import sys
import subprocess
import json
from pathlib import Path
from typing import Optional
import torch
import traceback
from transformers import GPT2LMHeadModel, GPT2Tokenizer

class Benchmark:
    def __init__(self, benchmark_config, model, tokenizer, checkpointing, output_dir, device=None):
        self.benchmark_config = benchmark_config

        self.model = model
        
        self.tokenizer = tokenizer
        self.checkpointing = checkpointing
        self.device = device if device is not None else torch.device("cuda" if torch.cuda.is_available() else "cpu")

        self.output_dir = Path(f"{benchmark_config.prefix if benchmark_config.prefix is not None else ''}{output_dir}")
        self.output_dir.mkdir(parents=True, exist_ok=True)

        self.lm_eval_dir = Path(__file__).resolve().parent / "lm-evaluation-harness"
        self.results_dir = self.output_dir / "results"
        self.summary_path = self.output_dir / "summary.json"
        
        self.checkpointing.load_model_states("recent")
        
        self._init_repo()

    def _init_repo(self):
        if not self.lm_eval_dir.exists():
            print("Cloning lm-evaluation-harness...")
            subprocess.run([
                "git", "clone", "https://github.com/EleutherAI/lm-evaluation-harness", str(self.lm_eval_dir)
            ], check=True)

            print("Installing lm-evaluation-harness...")
            subprocess.run(
                ["pip", "install", "-e", "."],
                cwd=str(self.lm_eval_dir),
                check=True
            )
        else:
            print("lm-evaluation-harness already installed.")
            

    def run_benchmarks(self):
        sys.path.insert(0, str(self.lm_eval_dir))

        from lm_eval.evaluator import simple_evaluate
        from .benchmark_wrapper import BenchmarkWrapper

        
        summary = self._load_existing_summary()
        status_report = {}

        model = self.model
        if self.benchmark_config.compile:
            print(f"Compiling model with mode: {self.benchmark_config.compile_mode}")
            model = torch.compile(model, mode=self.benchmark_config.compile_mode)

        wrapped_model = BenchmarkWrapper(
            model,
            self.tokenizer,
            device=self.device,
            batch_size=self.benchmark_config.batch_size,
            dtype=getattr(torch, self.benchmark_config.precision)
        )

        new_results = {}
        for task in self.benchmark_config.tasks:
            try:
                if task in summary:
                    print(f"Skipping {task}, already benchmarked.")
                    status_report[task] = {"status": "cached", "result": summary[task]}
                    continue

                print(f"Running benchmark for task: {task}")
                result = simple_evaluate(
                    model=wrapped_model,
                    tasks=[task],
                    num_fewshot=self.benchmark_config.num_fewshot,
                    batch_size=self.benchmark_config.batch_size,
                    limit=self.benchmark_config.limit if hasattr(self.benchmark_config, "limit") else None,
                    bootstrap_iters=self.benchmark_config.bootstrap_iters,
                    device=self.device,
                    log_samples=False
                )
                
                self.results_dir.mkdir(parents=True, exist_ok=True)
                
                def make_json_safe(obj):
                    if isinstance(obj, torch.device):
                        return str(obj)
                    elif isinstance(obj, dict):
                        return {k: make_json_safe(v) for k, v in obj.items()}
                    elif isinstance(obj, list):
                        return [make_json_safe(v) for v in obj]
                    elif isinstance(obj, tuple):
                        return tuple(make_json_safe(v) for v in obj)
                    else:
                        return obj
                
                with open(self.results_dir / f"{task}.json", "w") as f:
                    json.dump(make_json_safe(result), f, indent=2) 

                task_result = result["results"].get(task, {})
                summary[task] = task_result
                new_results[task] = task_result
                status_report[task] = {"status": "completed", "result": task_result}

            except Exception as e:
                print(f"⚠️  Warning: Benchmark for task '{task}' failed with error: {e}")
                traceback.print_exc()
                status_report[task] = {"status": "failed", "error": str(e)}

        self._save_summary(summary)
        self._print_status_report(status_report)

    def _load_existing_summary(self) -> dict:
        if self.summary_path.exists():
            with open(self.summary_path, "r") as f:
                return json.load(f)
        return {}

    def _save_summary(self, summary: dict):
        with open(self.summary_path, "w") as f:
            json.dump(summary, f, indent=2)

    def _print_status_report(self, report):
        def is_metric_key(k):
            return any(key in k.lower() for key in ["acc", "f1", "loss", "perplexity", "exact", "score"])

        def format_metric(k, v):
            
            clean_key = k.replace(",none", "")
            
            if "norm" in k:
                return None
            
            if isinstance(v, dict) and "value" in v:
                val = v["value"]
                stderr = v.get("stderr", None)
                if stderr is not None:
                    return f"{clean_key}: {val:.5f} ± {stderr:.5f}"
                else:
                    return f"{clean_key}: {val:.5f}"
            elif isinstance(v, (float, int)):
                return f"{clean_key}: {v:.5f}"
            return None

        print("\n=== Benchmark Summary Report ===")
        for task, info in report.items():
            status = info["status"]
            if status in {"completed", "cached"}:
                label = "✅" if status == "completed" else "📄"
                metrics = info["result"]
                filtered_metrics = {
                    k: v for k, v in metrics.items()
                    if is_metric_key(k)
                }
                formatted = [
                    format_metric(k, v)
                    for k, v in filtered_metrics.items()
                    if format_metric(k, v) is not None
                ]
                metric_str = ", ".join(formatted)
                print(f"{label} {task}: {status} | {metric_str}")
            elif status == "failed":
                print(f"❌ {task}: {status} | Error: {info['error']}")