#!/usr/bin/env python3
"""Batch Gemini archetype labeling (tier 0-5) — offline only."""

from __future__ import annotations

import argparse
import json
import logging
import sys
import time
from pathlib import Path

import polars as pl

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from app.candidates import load_candidates, load_candidates_list  # noqa: E402
from app.config import get_settings  # noqa: E402
from app.feature_store import load_silver_tiers  # noqa: E402
from app.gemini_client import (  # noqa: E402
    FLASH_MODEL_FALLBACKS,
    gemini_labels_enabled,
    generate_json,
    has_gemini_auth,
    load_prompt_template,
    resolve_flash_model,
)
from app.jd_requirements import JDRequirements  # noqa: E402
from app.logging_setup import configure_logging  # noqa: E402
from app.progress import ProgressTracker, format_duration  # noqa: E402

logger = logging.getLogger(__name__)

def _batch_size(settings) -> int:
    return int(getattr(settings, "gemini_label_batch_size", 50))


def _sleep_seconds(settings) -> float:
    return float(getattr(settings, "gemini_label_sleep_seconds", 0.2))


def _compact_candidate(candidate: dict) -> dict:
    profile = candidate.get("profile") or {}
    return {
        "candidate_id": candidate["candidate_id"],
        "title": profile.get("current_title"),
        "headline": profile.get("headline"),
        "summary": (profile.get("summary") or "")[:400],
        "yoe": profile.get("years_of_experience"),
        "country": profile.get("country"),
        "skills": [s.get("name") for s in (candidate.get("skills") or [])[:12]],
        "career_titles": [r.get("title") for r in (candidate.get("career_history") or [])[:4]],
    }


def label_batch_with_gemini(
    batch: list[dict],
    *,
    job_summary: str,
    settings,
) -> dict[str, int]:
    template = load_prompt_template("label_archetype.txt")
    prompt = (
        template.replace("{{JOB_SUMMARY}}", job_summary)
        .replace("{{CANDIDATES_JSON}}", json.dumps([_compact_candidate(c) for c in batch], indent=0))
    )
    payload = generate_json(
        prompt,
        settings=settings,
        model=resolve_flash_model(settings),
        temperature=0.0,
        model_fallbacks=FLASH_MODEL_FALLBACKS,
    )
    tiers: dict[str, int] = {}
    if isinstance(payload, list):
        for item in payload:
            if isinstance(item, dict) and "candidate_id" in item:
                tiers[str(item["candidate_id"])] = int(item.get("tier", 2))
    return tiers


def _rows_from_batch(
    batch: list[dict],
    gemini_tiers: dict[str, int],
    silver_tiers: dict[str, int],
) -> list[dict]:
    rows: list[dict] = []
    for c in batch:
        c_id = str(c["candidate_id"])
        rows.append(
            {
                "candidate_id": c_id,
                "gemini_tier": gemini_tiers.get(c_id, silver_tiers.get(c_id, 2)),
                "label_source": "gemini_flash" if c_id in gemini_tiers else "silver_fallback",
            }
        )
    return rows


def _silver_fallback_rows(batch: list[dict], silver_tiers: dict[str, int]) -> list[dict]:
    return [
        {
            "candidate_id": str(c["candidate_id"]),
            "gemini_tier": silver_tiers.get(str(c["candidate_id"]), 2),
            "label_source": "silver_fallback",
        }
        for c in batch
    ]


def build_gemini_tiers(
    candidates_path: Path,
    output_path: Path,
    *,
    settings,
    jd: JDRequirements,
    limit: int | None = None,
    force: bool = False,
    skip_api: bool = False,
) -> pl.DataFrame:
    if output_path.exists() and not force:
        logger.info("Loading cached Gemini tiers from %s", output_path)
        return pl.read_parquet(output_path)

    use_gemini = not skip_api and gemini_labels_enabled(settings) and has_gemini_auth(settings)
    if skip_api:
        logger.info("skip_api=True — writing silver-tier labels only (no Gemini spend)")
    elif not use_gemini:
        logger.warning("No Gemini credentials — writing silver-tier fallback labels")
    silver_path = settings.artifact_path("labels_silver.parquet")
    silver_tiers = load_silver_tiers(silver_path)

    job_summary = (
        f"Role: {jd.role_title}. Must-have: {', '.join(jd.must_have_skills[:8])}. "
        f"YOE: {jd.yoe_min}-{jd.yoe_max}. Country: {jd.preferred_country}."
    )

    rows: list[dict] = []
    batch: list[dict] = []
    candidates = load_candidates_list(candidates_path, limit=limit) if limit else list(load_candidates(candidates_path))
    total_candidates = len(candidates)
    batch_size = _batch_size(settings)
    sleep_seconds = _sleep_seconds(settings)
    total_batches = (total_candidates + batch_size - 1) // batch_size if use_gemini else 0

    if use_gemini:
        model = resolve_flash_model(settings)
        logger.info(
            "Gemini labeling: %d candidates in %d batches (batch_size=%d, model=%s)",
            total_candidates,
            total_batches,
            batch_size,
            model,
        )
        batch_progress = ProgressTracker(
            logger, label="Gemini labels", total=total_batches, log_every=1, unit="batches"
        )

        def flush_batch() -> None:
            nonlocal batch
            if not batch:
                return
            t0 = time.perf_counter()
            try:
                gemini_tiers = label_batch_with_gemini(batch, job_summary=job_summary, settings=settings)
                rows.extend(_rows_from_batch(batch, gemini_tiers, silver_tiers))
            except Exception as exc:
                logger.warning("Gemini batch failed (%s); using silver fallback", exc)
                rows.extend(_silver_fallback_rows(batch, silver_tiers))
            batch_progress.tick()
            logger.debug("Batch API call took %s", format_duration(time.perf_counter() - t0))
            batch = []
            time.sleep(sleep_seconds)

        for candidate in candidates:
            batch.append(candidate)
            if len(batch) >= batch_size:
                flush_batch()

        flush_batch()
        batch_progress.finish(message="labeling complete")
    else:
        candidate_progress = ProgressTracker(
            logger,
            label="Silver tier labels",
            total=total_candidates,
            log_every=max(total_candidates // 20, 1),
            unit="candidates",
        )
        for candidate in candidates:
            cid = str(candidate["candidate_id"])
            rows.append(
                {
                    "candidate_id": cid,
                    "gemini_tier": silver_tiers.get(cid, 2),
                    "label_source": "silver_fallback",
                }
            )
            candidate_progress.tick()
        candidate_progress.finish()

    frame = pl.DataFrame(rows)
    output_path.parent.mkdir(parents=True, exist_ok=True)
    frame.write_parquet(output_path)
    gemini_count = frame.filter(pl.col("label_source") == "gemini_flash").height
    logger.info(
        "Wrote Gemini tiers: %d rows (%d gemini_flash, %d silver_fallback) → %s",
        frame.height,
        gemini_count,
        frame.height - gemini_count,
        output_path,
    )
    return frame


def main() -> int:
    parser = argparse.ArgumentParser(description="Batch label candidate archetypes with Gemini Flash.")
    parser.add_argument("--candidates", type=Path, default=None)
    parser.add_argument("--out", type=Path, default=None)
    parser.add_argument("--limit", type=int, default=None, help="Limit candidates (testing).")
    parser.add_argument("--force", action="store_true")
    args = parser.parse_args()

    settings = get_settings()
    configure_logging(settings.log_level)
    settings.ensure_artifacts_dir()

    from app.jd_requirements import load_or_build_jd_requirements

    jd = load_or_build_jd_requirements(
        settings.job_description_path,
        settings.jd_requirements_path,
        settings=settings,
    )
    candidates_path = (args.candidates or settings.candidates_path).resolve()
    out_path = (args.out or settings.artifact_path("gemini_tiers.parquet")).resolve()

    if not has_gemini_auth(settings):
        logger.warning("No Gemini credentials — writing silver-tier fallback labels")

    build_gemini_tiers(
        candidates_path,
        out_path,
        settings=settings,
        jd=jd,
        limit=args.limit,
        force=args.force,
    )
    return 0


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