from argparse import ArgumentParser, ArgumentDefaultsHelpFormatter
from pathlib import Path
from dataclasses import dataclass
import resource
import docker
import time
import re
import warnings

from gso.utils.io import load_gso_dataset, load_gso_predictions
from gso.harness.environment.docker_utils import cleanup_docker
from gso.harness.grading.utils import get_dataset_from_preds
from gso.harness.grading.grade import grade_instance
from gso.harness.grading.report import make_run_report
from gso.harness.utils import retag_remote_to_local_image
from gso.utils.multiprocess import run_tasks_in_parallel_iter
from gso.constants import HIGH_RESOURCE_REPOS
from gso.data.dataset import GSOInstance

warnings.filterwarnings("ignore", category=RuntimeWarning, module="numpy")


@dataclass
class GradeInstanceTask:
    instance: GSOInstance
    pred: dict
    rm_image: bool
    run_id: str
    timeout: int | None = None
    resource_level: str = "low"
    reformat_reports: bool = False


def grad_instance_mp(task: GradeInstanceTask):
    max_retries = 3
    sleep_time = 330
    instance_id = task.instance.instance_id

    # if reformat_reports, set max_retries to 0
    if task.reformat_reports:
        max_retries = 0

    for attempt in range(max_retries + 1):
        result = grade_instance(
            task.instance,
            task.pred,
            task.rm_image,
            task.run_id,
            task.timeout,
            task.reformat_reports,
            retry_count=attempt,
            max_retries=max_retries,
        )

        if result is None:
            return result

        result_id, report, test_output_path = result

        # read test_output
        with open(test_output_path, "r") as f:
            test_output = f.read()

        # if base_successfully_run is False, retry
        if (
            report
            and instance_id in report
            and not report[instance_id].get("base_successfully_run", False)
        ):
            if attempt < max_retries:
                time.sleep(sleep_time)
                continue

        # if HTTP error:
        if re.search(r"HTTP.*Error|Error.*HTTP", test_output, re.IGNORECASE):
            if attempt < max_retries:
                time.sleep(sleep_time)
                continue

        return result


def run_instances(
    predictions: dict,
    instances: list[GSOInstance],
    max_workers: int,
    run_id: str,
    timeout: int,
    reformat_reports: bool = False,
):
    """
    Run all instances for the given predictions in parallel.

    Args:
        predictions (dict): Predictions dict generated by the model
        instances (list): List of instances
        max_workers (int): Maximum number of workers
        run_id (str): Run ID
        timeout (int): Timeout for running tests
        reformat_reports (bool): Whether to rewrite reports for instances that have already been run
    """
    client = docker.from_env()

    # print number of existing instance images
    instance_image_ids = {x.instance_image_key for x in instances}
    retag_remote_to_local_image(instances, client)
    existing_images = {
        tag
        for i in client.images.list(all=True)
        for tag in i.tags
        if tag in instance_image_ids
    }
    print(f"Found {len(existing_images)} existing instance images.")

    # warn if any instance images are missing
    # TODO: pull images as and when needed
    missing_images = instance_image_ids - existing_images
    if missing_images:
        print(f"{"="*10} Warning: {len(missing_images)} images missing {"="*10}")
        print(f"Missing: {missing_images}\n")
        print("Please pull/build missing images before running the instances.")
        print(f"{"="*10} End {"="*10}\n")
        return

    grading_tasks = []
    for instance in instances:
        is_high_resource = instance.repo in HIGH_RESOURCE_REPOS
        timeout = timeout * 2 if is_high_resource else timeout
        git = GradeInstanceTask(
            instance=instance,
            pred=predictions[instance.instance_id],
            rm_image=False,
            run_id=run_id,
            timeout=timeout,
            reformat_reports=reformat_reports,
            resource_level="high" if is_high_resource else "low",
        )
        grading_tasks.append(git)

    # split tasks into two groups based on resource level
    lr_tasks = [t for t in grading_tasks if t.resource_level == "low"]
    hr_tasks = [t for t in grading_tasks if t.resource_level == "high"]
    successful, failed = [], []

    if lr_tasks:
        results = run_tasks_in_parallel_iter(
            grad_instance_mp,
            tasks=lr_tasks,
            num_workers=max_workers,
            timeout_per_task=3600,
            use_progress_bar=True,
            progress_bar_desc="Grading instances [type: lr]",
        )

        for task, config in zip(results, lr_tasks):
            if task.is_success():
                successful.append(config.instance.instance_id)
            else:
                failed.append(config.instance.instance_id)
                if task.is_timeout():
                    print(f"Timeout: {config.instance.instance_id}")
                elif task.is_process_expired():
                    print(f"Process expired: {config.instance.instance_id}")
                elif task.is_exception():
                    print(f"Error: {config.instance.instance_id}")
                    print(task.exception_tb)

    # clean up all containers
    cleanup_docker(client)

    if hr_tasks:
        results = run_tasks_in_parallel_iter(
            grad_instance_mp,
            tasks=hr_tasks,
            num_workers=1,  # force sequential execution
            timeout_per_task=5400,  # increase timeout
            use_progress_bar=True,
            progress_bar_desc="Grading instances [type: hr]",
        )

        for task, config in zip(results, hr_tasks):
            if task.is_success():
                successful.append(config.instance.instance_id)
            else:
                failed.append(config.instance.instance_id)
                if task.is_timeout():
                    print(f"Timeout: {config.instance.instance_id}")
                elif task.is_process_expired():
                    print(f"Process expired: {config.instance.instance_id}")
                elif task.is_exception():
                    print(f"Error: {config.instance.instance_id}")
                    print(task.exception_tb)

    if len(failed) == 0:
        print("All instances ran successfully.")
    else:
        print(f"{len(failed)} instance images failed to build.")

    # clean up all containers
    cleanup_docker(client)


def main(
    dataset_name,
    split,
    instance_ids,
    predictions_path,
    max_workers,
    open_file_limit,
    timeout,
    run_id,
    reformat_reports,
    rerun_all=False,
    report_dir=".",
    verbose=False,
):
    """Run eval harness for a given dataset and predictions."""
    assert len(run_id) > 0, "Run ID must be provided"
    if report_dir is not None:
        report_dir = Path(report_dir)
        if not report_dir.exists():
            report_dir.mkdir(parents=True)

    # load the prediction: {instance_id: prediction}
    predictions = load_gso_predictions(predictions_path, dataset_name, split)
    predictions = {p["instance_id"]: p for p in predictions}

    # filter dataset to predictions
    predictions, dataset = get_dataset_from_preds(
        dataset_name,
        split,
        instance_ids,
        predictions,
        run_id,
        reformat_reports,
        exclude_completed=not rerun_all,
    )
    full_dataset = load_gso_dataset(dataset_name, split)

    # Set open file limit
    resource.setrlimit(resource.RLIMIT_NOFILE, (open_file_limit, open_file_limit))
    client = docker.from_env()

    print(f"Running {len(dataset)} unevaluated instances...")
    if not dataset:
        print("No instances to run.")
    else:
        run_instances(
            predictions,
            dataset,
            max_workers,
            run_id,
            timeout,
            reformat_reports=reformat_reports,
        )

    return make_run_report(
        predictions, full_dataset, run_id, client=None, verbose=verbose
    )


if __name__ == "__main__":
    parser = ArgumentParser(
        description="Run eval harness for a given dataset and predictions.",
        formatter_class=ArgumentDefaultsHelpFormatter,
    )

    # common args
    parser.add_argument(
        "--run_id", type=str, required=True, help="Run ID - identifies the run"
    )
    parser.add_argument(
        "--dataset_name",
        type=str,
        default="gso-bench/gso",
        help="Name of HF dataset to use or local json/jsonl file",
    )

    parser.add_argument("--split", type=str, default="test", help="Split to use")
    parser.add_argument(
        "--instance_ids",
        nargs="+",
        type=str,
        help="Instance IDs to run (space separated)",
    )

    parser.add_argument(
        "--predictions_path", type=str, help="Path to predictions file.", required=True
    )

    # execution args
    parser.add_argument(
        "--report_dir",
        type=str,
        default=".",
        help="Directory to store the reports",
    )
    parser.add_argument(
        "--reformat_reports",
        action="store_true",
        help="Reformat and rewrite reports for instances that have already been run",
    )
    parser.add_argument(
        "--rerun_all",
        action="store_true",
        help="Rerun all instances, even if they have already been run. Default: False",
    )
    parser.add_argument(
        "--max_workers",
        type=int,
        default=10,
        help="Max workers for parallel processing",
    )
    parser.add_argument(
        "--open_file_limit", type=int, default=8192, help="Open file limit"
    )
    parser.add_argument(
        "--timeout",
        type=int,
        default=1_800,
        help="Timeout (in seconds) for running an eval for each instance",
    )
    parser.add_argument(
        "--verbose",
        action="store_true",
        help="Print detailed per-run summaries",
    )

    args = parser.parse_args()
    main(**vars(args))
