from __future__ import annotations

import argparse
import json
import os
from pathlib import Path

import modal

from dotenv import load_dotenv
load_dotenv()

def _require_env(name: str) -> str:
    value = os.environ.get(name, "").strip()
    if not value:
        raise ValueError(f"Missing required env var: {name}")
    return value


def _seed_dummy_dataset(app_name: str) -> None:
    seed_fn = modal.Function.from_name(app_name, "seed_dummy_dataset")
    seed = seed_fn.remote(num_images=6, image_size=224, dataset_name="dummy")
    print(f"Seeded dataset at {seed['dataset_root']}")


def main(
    models: list[str], 
    submitter: str,
    stimuli: str | None, 
    use_s3: bool = True, 
    skip_seed: bool = False
) -> None:

    _require_env("HACKATHON_MODAL_ENABLE")
    _require_env("HACKATHON_MODEL_REGISTRY")

    # Use custom stimuli if provided, otherwise use env var
    if stimuli:
        stimuli_path = Path(stimuli).resolve()
        if not stimuli_path.exists():
            raise ValueError(f"Stimuli file not found: {stimuli_path}")
        os.environ["HACKATHON_STIMULI_CATALOG"] = str(stimuli_path)
        print(f"Using custom stimuli: {stimuli_path}")
    else:
        _require_env("HACKATHON_STIMULI_CATALOG")

    # Set S3 mode if requested
    if use_s3:
        os.environ["HACKATHON_USE_S3"] = "true"
        print("S3 mode enabled - will use extract_embeddings_s3 function")

    app_name = os.environ.get("HACKATHON_MODAL_APP", "iclr2026-eval")

    # Only seed dummy dataset if not using real images
    if not skip_seed and not use_s3:
        _seed_dummy_dataset(app_name)
    elif skip_seed:
        print("Skipping dummy dataset seeding")
    elif use_s3:
        print("S3 mode: skipping dummy dataset seeding (using real images)")

    from app import submit_blue

    payload = json.dumps({"models": models})
    msg, leaderboard, pairwise = submit_blue(submitter, payload)
    print(f"Submission message: {msg}")
    print(f"Submission leaderboard: {leaderboard.tail(1).to_dict(orient='records')}")
    print(f"Submission pairwise: {pairwise.to_dict(orient='records')}")
    assert not pairwise.empty, "Pairwise table should not be empty."

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Blue family smoke test")
    parser.add_argument(
        "--models",
        type=str,
        nargs="+",
        default=None,
        help="List of models to submit. If not provided, uses predefined model sets.",
    )
    parser.add_argument(
        "--submitter",
        type=str,
        default="test",
        help="Submitter name to use for the submission.",
    )
    parser.add_argument(
        "--stimuli",
        type=str,
        default=None,
        help="Path to custom stimuli JSONL file. If not provided, uses HACKATHON_STIMULI_CATALOG env var.",
    )
    parser.add_argument(
        "--s3",
        action="store_true",
        help="Use S3-backed datasets (requires aws-s3-credentials Modal secret).",
    )
    parser.add_argument(
        "--skip-seed",
        action="store_true",
        help="Skip seeding dummy dataset (use when testing with real images).",
    )
    args = parser.parse_args()

    models = args.models if args.models is not None else ["resnet18", "resnet34"]
    submitter = args.submitter

    main(models, submitter, args.stimuli, args.s3, args.skip_seed)