"""
Part 2 entry point: run race classification, skin tone extraction, and
CLIP/face embeddings over every image in a Part 1 manifest, one image at
a time (resumable, same philosophy as Part 1's run_generation).

Which analyses run is selectable via `components` -- e.g. run only skin
tone + face embeddings and skip CLIP (useful when CLIP's ~4GB download
won't fit on disk). Default is all four.

Usage:
    import pandas as pd
    from src.analyze import run_analysis

    manifest_df = pd.read_parquet("runs/my_audit/manifest.parquet")
    analysis_path, enriched_df = run_analysis(manifest_df, output_dir="runs/my_audit")
    # or a subset:
    run_analysis(manifest_df, "runs/my_audit", components={"skin", "face"})
"""
import glob
import os

from PIL import Image

from .zerogpu import GPU
from .analysis_writer import AnalysisWriter
from .race_classification import classify_race
from .skin_tone import analyze_skin_tone_from_img
from .embeddings import embed_clip, embed_face

ALL_COMPONENTS = {"race", "skin", "clip", "face"}


@GPU(duration=15)
def _zerogpu_startup_probe():
    """ZeroGPU requires at least one @spaces.GPU function to exist at startup.
    We deliberately DON'T route the real analysis through a per-image GPU call:
    ZeroGPU forks a fresh process per call, so lazily-loaded models (the 4GB
    CLIP + InsightFace) reloaded on every image and blew past the GPU token's
    lifetime ("expired GPU proxy token"). Instead the model steps run in the
    main process (below), where models load once and are reused. This probe
    exists only to satisfy the startup requirement."""
    return True


def _run_model_analyses(pil_image, components):
    """Skin tone (BiSeNet), CLIP, and face embeddings for one image. Runs in
    the main process (CPU on ZeroGPU) so the heavy models load once and are
    reused across images. Race classification is separate -- it's a Gemini API
    call handled in _analyze_one.

    Collects step failures in fields["_errors"] (a list) so a bad step doesn't
    lose the others; _analyze_one folds those into the final `error` column."""
    fields = {"_errors": []}

    if "skin" in components:
        skin = analyze_skin_tone_from_img(pil_image)
        if "error" in skin:
            fields["skin_tone_error"] = skin["error"]
        else:
            fields["L"], fields["a"], fields["b"], fields["h_star"] = (
                skin["L"], skin["a"], skin["b"], skin["h_star"]
            )

    if "clip" in components:
        try:
            fields["clip_emb"] = embed_clip(pil_image).tolist()
        except Exception as e:
            fields["clip_emb"] = None
            fields["_errors"].append(f"clip:{type(e).__name__}:{e}")

    if "face" in components:
        try:
            face_vec, det_score, bbox, n_faces, status = embed_face(pil_image)
            fields["face_emb"] = face_vec.tolist() if face_vec is not None else None
            fields["face_det_score"] = det_score
            fields["face_bbox_xywh"] = bbox
            fields["n_faces_detected"] = n_faces
            fields["embed_status"] = status
        except Exception as e:
            fields["embed_status"] = f"error:{type(e).__name__}"
            fields["_errors"].append(f"face:{type(e).__name__}:{e}")

    return fields


def _analyze_one(pil_image, components, gemini_api_key=None):
    """Runs the selected analyses for one image. Never raises -- a failure
    of an individual step is recorded in the `error` field so one bad step
    doesn't lose the others. GPU-bound steps run in _gpu_analyze_one (which
    ZeroGPU attaches a GPU to); the Gemini race call stays on CPU.

    gemini_api_key: passed to the classifier so callers can supply a
    per-request key (bring-your-own-key UI) instead of relying on the
    GOOGLE_API_KEY env var."""
    fields = {}
    errors = []

    if "race" in components:
        try:
            fields["race_classification"] = classify_race(pil_image, api_key=gemini_api_key)
        except Exception as e:
            fields["race_classification"] = None
            errors.append(f"race_classification:{type(e).__name__}:{e}")

    model_fields = _run_model_analyses(pil_image, components)
    errors.extend(model_fields.pop("_errors", []))
    fields.update(model_fields)

    if errors:
        fields["error"] = " | ".join(errors)
    return fields


def _clear_prior_analysis(output_dir):
    """Remove any existing analysis chunks/outputs so a run starts fresh.
    Used when `force=True` (e.g. the user changed which components to run --
    resumability is image-level, not component-level, so re-running with a
    different component set requires starting over)."""
    chunk_dir = os.path.join(output_dir, "analysis_chunks")
    for f in glob.glob(os.path.join(chunk_dir, "chunk_*.parquet")):
        os.remove(f)
    for name in ("analysis.parquet", "analyzed_manifest.parquet"):
        p = os.path.join(output_dir, name)
        if os.path.exists(p):
            os.remove(p)


def run_analysis(manifest_df, output_dir, components=None, force=False,
                 on_progress=None, gemini_api_key=None):
    """
    manifest_df: the Part 1 manifest (must have image_id, file_path columns)
    output_dir: same run directory used for generation -- analysis_chunks/
        and analysis.parquet are written alongside images/ and manifest.parquet
    components: which analyses to run, subset of {"race","skin","clip","face"}.
        None => all four.
    force: if True, discard any prior analysis for this run and start fresh.
        (Resumability is image-level; re-running with a different component
        set needs force=True to actually re-analyze already-seen images.)
    on_progress: optional callback(done, total)
    gemini_api_key: per-request Google API key for race classification
        (bring-your-own-key); falls back to the GOOGLE_API_KEY env var.

    Returns: (analysis_path, enriched_df) where enriched_df is manifest_df
    left-joined with the analysis results on image_id.
    """
    components = ALL_COMPONENTS if components is None else set(components)
    if not components:
        raise ValueError("No analysis components selected.")

    if force:
        _clear_prior_analysis(output_dir)

    writer = AnalysisWriter(output_dir)

    total = len(manifest_df)
    done = 0

    for _, row in manifest_df.iterrows():
        image_id = row["image_id"]

        if writer.is_done(image_id):
            done += 1
            if on_progress:
                on_progress(done, total)
            continue

        img = Image.open(row["file_path"]).convert("RGB")
        fields = _analyze_one(img, components, gemini_api_key=gemini_api_key)
        writer.add_row(image_id, **fields)

        done += 1
        if on_progress:
            on_progress(done, total)

    analysis_path, analysis_df = writer.finalize()
    enriched_df = manifest_df.merge(analysis_df, on="image_id", how="left")
    enriched_path = os.path.join(output_dir, "analyzed_manifest.parquet")
    enriched_df.to_parquet(enriched_path, index=False)
    return enriched_path, enriched_df