from __future__ import annotations

import argparse
from collections.abc import Sequence

from google import genai
from google.genai import types

from copienator import configuration as config
from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    atomic_write_json,
    evaluation_parser,
    execute,
    read_json,
    workspace_from_args,
)


def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
    if client is None:
        if not config.API_KEY:
            raise CliError("GEMINI_API_KEY is not configured")
        client = genai.Client(api_key=config.API_KEY)
    batches = (
        (
            "flash",
            workspace.root / "batch_requests_flash.jsonl",
            config.MODEL_FLASH_ID,
            f"flash-correction-{workspace.name}",
        ),
        (
            "pro",
            workspace.root / "batch_requests_pro.jsonl",
            config.MODEL_PRO_ID,
            f"pro-correction-{workspace.name}",
        ),
    )
    manifest = {
        "version": 1,
        "evaluation": workspace.name,
        "jobs": {},
    }
    if workspace.batch_jobs_file.is_file():
        previous = read_json(workspace.batch_jobs_file)
        if isinstance(previous, dict) and isinstance(previous.get("jobs"), dict):
            manifest["jobs"] = previous["jobs"]
    started = 0
    for tier, file_path, model_id, display_name in batches:
        if not file_path.is_file():
            print(f"Skipping {model_id}: {file_path.name} does not exist.")
            continue
        if file_path.stat().st_size == 0:
            print(f"Skipping {model_id}: {file_path.name} is empty.")
            continue
        print(f"Uploading {file_path.name} for model {model_id}...")
        uploaded = client.files.upload(
            file=str(file_path),
            config=types.UploadFileConfig(
                display_name=f"{display_name}-input",
                mime_type="jsonl",
            ),
        )
        job = client.batches.create(
            model=model_id,
            src=uploaded.name,
            config={"display_name": display_name},
        )
        started += 1
        manifest["jobs"][tier] = {
            "name": job.name,
            "display_name": display_name,
            "model": model_id,
            "request_file": file_path.name,
        }
        atomic_write_json(workspace.batch_jobs_file, manifest)
        print(f"Started batch job: {job.name}")
    if not started:
        print("No non-empty batch request files were found.")
        return ExitCode.PARTIAL
    return ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    return evaluation_parser("Upload correction JSONL files and start Gemini batches")


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()
    return execute(parser, argv, lambda args: run(workspace_from_args(args)))


if __name__ == "__main__":
    raise SystemExit(main())