from __future__ import annotations

import argparse
from collections.abc import Sequence

from google import genai

import config
from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    atomic_write_bytes,
    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)
    matching = []
    if workspace.batch_jobs_file.is_file():
        manifest = read_json(workspace.batch_jobs_file)
        jobs = manifest.get("jobs") if isinstance(manifest, dict) else None
        if not isinstance(jobs, dict):
            raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
        matching = [
            client.batches.get(name=entry["name"])
            for entry in jobs.values()
            if isinstance(entry, dict) and isinstance(entry.get("name"), str)
        ]
    else:
        matching = [
            job
            for job in client.batches.list()
            if workspace.name in str(getattr(job, "display_name", ""))
        ]
    if not matching:
        raise CliError(
            f"No batch jobs found for evaluation {workspace.name!r}"
        )
    for job in matching:
        state = job.state.name if hasattr(job.state, "name") else job.state
        print(f"{job.display_name}: {state}")
        if state != "JOB_STATE_SUCCEEDED":
            print("Not all matching jobs have succeeded yet.")
            return ExitCode.PARTIAL

    chunks = []
    incomplete = False
    for job in matching:
        destination = getattr(job, "dest", None)
        file_name = getattr(destination, "file_name", None)
        if not file_name:
            print(f"Warning: {job.display_name} has no output file.")
            incomplete = True
            continue
        payload = client.files.download(file=file_name)
        chunks.append(payload.rstrip(b"\n"))
    if not chunks:
        return ExitCode.PARTIAL
    output_path = workspace.batched_correction_result_file
    atomic_write_bytes(output_path, b"\n".join(chunks) + b"\n")
    print(f"Saved combined batch results to {output_path}")
    return ExitCode.PARTIAL if incomplete else ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    return evaluation_parser("Download and combine correction batch results")


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())