from __future__ import annotations

import argparse
from collections.abc import Sequence
from pathlib import Path

from google import genai

from copienator import configuration as config
from copienator import (
    CliError,
    EvaluationWorkspace,
    ExitCode,
    atomic_write_bytes,
    execute,
    read_json,
    standard_parser,
    workspace_from_args,
)


def _client():
    if not config.API_KEY:
        raise CliError("GEMINI_API_KEY is not configured")
    return genai.Client(api_key=config.API_KEY)


def list_jobs(*, client=None) -> ExitCode:
    client = client or _client()
    print("Fetching recent batch jobs...")
    jobs = list(client.batches.list())
    for job in jobs:
        state = job.state.name if hasattr(job.state, "name") else job.state
        print(f"{job.name}: {state}")
        if getattr(job, "display_name", None):
            print(f"  Display name: {job.display_name}")
        if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
            print(f"  Error: {job.error}")
        destination = getattr(job, "dest", None)
        if state == "JOB_STATE_SUCCEEDED" and getattr(
            destination, "file_name", None
        ):
            print(f"  Output file: {destination.file_name}")
    if not jobs:
        print("No batch jobs found.")
    return ExitCode.SUCCESS


def check_evaluation_jobs(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
    """Only report readiness after checking every job recorded for this evaluation."""
    if not workspace.batch_jobs_file.is_file():
        print("Impossible de vérifier les batchs : batch_jobs.json est absent.")
        return ExitCode.PARTIAL
    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}")
    if not jobs:
        print("Aucun job enregistré pour cette évaluation.")
        return ExitCode.PARTIAL
    if any(not isinstance(entry, dict) or not isinstance(entry.get("name"), str)
           or not entry["name"].strip() for entry in jobs.values()):
        raise CliError(f"Invalid batch job in {workspace.batch_jobs_file}")
    client = client or _client()
    ready = True
    for tier, entry in jobs.items():
        job = client.batches.get(name=entry["name"])
        state = job.state.name if hasattr(job.state, "name") else job.state
        print(f"{tier} — {entry['name']}: {state}")
        if state != "JOB_STATE_SUCCEEDED":
            ready = False
            if getattr(job, "error", None):
                print(f"  Erreur : {job.error}")
        elif not getattr(getattr(job, "dest", None), "file_name", None):
            ready = False
            print("  Le fichier de résultats n’est pas encore disponible.")
    if ready:
        print("Tous les batchs de l’évaluation ont réussi. Les résultats sont prêts à récupérer.")
        return ExitCode.SUCCESS
    print("Les résultats ne sont pas tous prêts. Consultez à nouveau cette étape plus tard.")
    return ExitCode.PARTIAL


def download_job(
    job_name: str,
    *,
    output: Path | None = None,
    client=None,
) -> ExitCode:
    client = client or _client()
    job = client.batches.get(name=job_name)
    state = job.state.name if hasattr(job.state, "name") else job.state
    print(f"State: {state}")
    if state != "JOB_STATE_SUCCEEDED":
        if state == "JOB_STATE_FAILED" and getattr(job, "error", None):
            print(f"Error: {job.error}")
        return ExitCode.PARTIAL
    destination = getattr(job, "dest", None)
    file_name = getattr(destination, "file_name", None)
    if not file_name:
        print("Job succeeded but no output file was found.")
        return ExitCode.PARTIAL
    payload = client.files.download(file=file_name)
    output_path = output or Path(f"results_{job_name.replace('/', '_')}.jsonl")
    atomic_write_bytes(output_path, payload)
    print(f"Saved batch results to {output_path}")
    return ExitCode.SUCCESS


def build_parser() -> argparse.ArgumentParser:
    parser = standard_parser("List or download Gemini correction batch jobs")
    parser.add_argument("--download", metavar="JOB_NAME")
    parser.add_argument("--output", type=Path, help="Downloaded JSONL destination")
    parser.add_argument("--evaluation", type=Path,
                        help="Check readiness of jobs recorded in this evaluation's batch_jobs.json")
    return parser


def main(argv: Sequence[str] | None = None) -> int:
    parser = build_parser()

    def handle(args: argparse.Namespace) -> ExitCode:
        if args.evaluation is not None:
            if args.download or args.output is not None:
                raise CliError("--evaluation cannot be combined with --download or --output",
                               ExitCode.INVALID_ARGUMENTS)
            return check_evaluation_jobs(workspace_from_args(args))
        if args.output is not None and not args.download:
            raise CliError("--output requires --download", ExitCode.INVALID_ARGUMENTS)
        if args.download:
            return download_job(args.download, output=args.output)
        return list_jobs()

    return execute(parser, argv, handle)


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