#!/usr/bin/env python
import argparse
import copy
import json
import os
from concurrent.futures import ThreadPoolExecutor
from datetime import datetime

from dotenv import load_dotenv
from smolagents import InferenceClientModel, LiteLLMModel, OpenAIServerModel
from tqdm import tqdm

from screensuite import registry
from screensuite.basebenchmark import EvaluationConfig

RESULTS_DIR = "./output"
if not os.path.exists(RESULTS_DIR):
    RESULTS_DIR = os.path.join(os.path.dirname(__file__), "results")
    os.makedirs(RESULTS_DIR, mode=0o775, exist_ok=True)

load_dotenv()


def launch_test(model, benchmarks, original_evaluation_config):
    evaluation_config = copy.deepcopy(original_evaluation_config)  # NOTE: important!
    if model.model_id is None:
        model_name = "custom-endpoint"
    elif model.model_id.startswith("http"):
        model_name = model.model_id.split("/")[-1][:10]
    else:
        model_name = model.model_id

    if evaluation_config.run_name is None:
        evaluation_config.run_name = f"{model_name.replace('/', '-')}_{datetime.now().strftime('%Y-%m-%d')}"

    print(f"===== Running evaluation under name: {evaluation_config.run_name} =====")

    # Load already processed benchmarks from results.jsonl
    output_results_file = f"{RESULTS_DIR}/{evaluation_config.run_name}.jsonl"
    processed_benchmarks = set()
    try:
        with open(output_results_file, "r") as f:
            for line in f:
                try:
                    data = json.loads(line)
                    if "benchmark_name" in data:
                        processed_benchmarks.add(data["benchmark_name"])
                except json.JSONDecodeError:
                    continue
    except FileNotFoundError:
        pass

    print("-> Found these processed benchmarks: ", processed_benchmarks)

    for benchmark in tqdm(sorted(benchmarks, key=lambda b: b.name), desc="Running benchmarks"):
        if benchmark.name in processed_benchmarks:
            print(f"Skipping already processed benchmark: {benchmark.name}")
            continue
        if "multistep" not in benchmark.tags:
            print("=" * 100)
            print(f"Running benchmark: {benchmark.name}")
            try:
                benchmark.load()

                results = benchmark.evaluate(
                    model,
                    evaluation_config,
                )
                print(f"Results for {benchmark.name}: {results}")

                # Save metrics to JSONL file
                metrics_entry = {"benchmark_name": benchmark.name, "metrics": results._metrics}
                with open(output_results_file, "a") as f:
                    f.write(json.dumps(metrics_entry) + "\n")
            except Exception as e:
                print(f"Error running benchmark {benchmark.name}: {e}")
                continue


def main():
    parser = argparse.ArgumentParser(description="Run benchmarks with optional run name")
    parser.add_argument("--run-name", type=str, help="Name of the run to continue or create", default=None)
    parser.add_argument(
        "--tag",
        type=str,
        nargs="+",
        help="Tags to filter benchmarks (can provide multiple tags)",
        default=["to_evaluate"],
    )
    parser.add_argument("--n-samples", type=int, help="Number of samples", default=300)
    args = parser.parse_args()

    # Get all registered benchmarks
    if not args.tag:  # This handles both empty list and None
        all_benchmarks = registry.list_all()
    else:
        all_benchmarks = registry.get_by_tags(tags=args.tag, match_all=False)

    non_multistep_benchmarks = [benchmark for benchmark in all_benchmarks if "multistep" not in benchmark.tags]
    print(f"Selected {len(non_multistep_benchmarks)} benchmarks:")
    for benchmark in non_multistep_benchmarks:
        print(f"- {benchmark.name}")

    # assert len(non_multistep_benchmarks) == 9

    evaluation_config = EvaluationConfig(
        parallel_workers=6,
        max_samples_to_test=args.n_samples,
        run_name=args.run_name,
    )
    evaluation_config_moreworkers = EvaluationConfig(
        parallel_workers=32,
        max_samples_to_test=args.n_samples,
        run_name=args.run_name,
    )
    model_qwenvl_72b_custom = InferenceClientModel(
        model_id="https://n5wr7lfx6wp94tvl.us-east-1.aws.endpoints.huggingface.cloud",
        # model_id="Qwen/Qwen2.5-VL-72B-Instruct",
        # provider="nebius",
        max_tokens=4096,
    )

    model_openai = OpenAIServerModel(
        model_id="gpt-4o",
        max_tokens=4096,
    )

    model_qwenvl_32b = InferenceClientModel(
        model_id="Qwen/Qwen2.5-VL-32B-Instruct",
        provider="fireworks-ai",
        max_tokens=4096,
    )

    model_qwenvl_72b = InferenceClientModel(
        model_id="Qwen/Qwen2.5-VL-72B-Instruct",
        provider="nebius",
        max_tokens=4096,
    )

    model_uitars_15_custom = InferenceClientModel(
        model_id="https://a1ktmpegjfwv9n2y.us-east-1.aws.endpoints.huggingface.cloud",
        max_tokens=4096,
    )

    evaluation_config_uitars = EvaluationConfig(
        parallel_workers=16,
        max_samples_to_test=args.n_samples,
        run_name=args.run_name,
        use_uitars_action_space=True,
        image_before_text=False,
    )

    model_h_custom = InferenceClientModel(
        model_id="https://iw3ycz1skszlusk8.us-east-1.aws.endpoints.huggingface.cloud",
        max_tokens=4096,
    )

    evaluation_config_h = EvaluationConfig(
        parallel_workers=16,
        max_samples_to_test=args.n_samples,
        run_name=args.run_name,
        use_h_action_space=True,
        image_before_text=True,
    )
    assert evaluation_config_h.custom_system_prompt is not None

    model_claude = LiteLLMModel(
        model_id="anthropic/claude-sonnet-4-20250514",
        max_tokens=4096,
    )

    evaluation_couples = [
        # (evaluation_config, model_qwenvl_72b_custom),
        # (evaluation_config, model_qwenvl_72b),
        (evaluation_config_moreworkers, model_openai),
        (evaluation_config, model_qwenvl_32b),
        (evaluation_config_uitars, model_uitars_15_custom),
        (evaluation_config_h, model_h_custom),
        # (evaluation_config, model_claude),
    ]

    with ThreadPoolExecutor(max_workers=6) as executor:
        futures = [
            executor.submit(launch_test, model, non_multistep_benchmarks, evaluation_config)
            for evaluation_config, model in evaluation_couples
        ]
        for future in futures:
            future.result()

    # launch_test(model_qwenvl_32b, non_multistep_benchmarks, evaluation_config)


if __name__ == "__main__":
    main()
