"""Composants KFP v2 du pipeline d'évaluation de prompts.

Chaque fonction est un composant Python « léger » : Vertex AI l'exécute dans un
conteneur `python:3.11-slim`. Elle doit donc être autonome (imports DANS la
fonction, aucune référence au reste du module).
Les tests unitaires appellent `<composant>.python_func(...)` directement.
"""
from typing import NamedTuple

from kfp import dsl
from kfp.dsl import Artifact, ClassificationMetrics, Dataset, Input, Metrics, Output

BASE_IMAGE = "python:3.11-slim"
GENAI = "google-genai==2.31.0"


@dsl.component(base_image=BASE_IMAGE)
def prepare_dataset(
    source_path: str,
    source_sha: str,
    labels: list,
    dataset: Output[Dataset],
) -> NamedTuple("PrepOutputs", [("n_examples", int), ("fingerprint", str)]):
    """Valide le jeu de référence, retire les doublons, calcule une empreinte.

    `source_sha` (empreinte du fichier brut calculée à la soumission) entre dans la clé
    de cache : sans lui, un fichier modifié au même chemin renverrait l'ancien résultat.
    """
    import hashlib
    import json
    import logging
    from collections import Counter

    log = logging.getLogger("prepare_dataset")
    with open(source_path, "rb") as f:
        raw = f.read()
    actual = hashlib.sha256(raw).hexdigest()[:16]
    if source_sha and actual != source_sha:
        raise ValueError(f"le fichier lu ({actual}) n'est pas celui soumis ({source_sha})")
    rows, seen, rejected = [], set(), 0
    for n, line in enumerate(raw.decode("utf-8").splitlines(), 1):
        if not line.strip():
            continue
        row = json.loads(line)
        text = " ".join(str(row.get("text", "")).split())
        key = text.casefold()
        if not text or row.get("label") not in labels:
            log.warning("ligne %d rejetée : texte vide ou label inconnu", n)
            rejected += 1
            continue
        if key in seen:
            log.warning("ligne %d rejetée : doublon de %r", n, text[:40])
            rejected += 1
            continue
        seen.add(key)
        rows.append({"id": row["id"], "text": text, "label": row["label"]})
    if len(rows) < 5:
        raise ValueError(f"jeu de référence trop petit : {len(rows)} exemples valides")

    payload = "\n".join(json.dumps(r, ensure_ascii=False, sort_keys=True) for r in rows)
    fingerprint = hashlib.sha256(payload.encode()).hexdigest()[:16]
    with open(dataset.path, "w", encoding="utf-8") as f:
        f.write(payload + "\n")
    dataset.metadata.update(
        n_examples=len(rows),
        rejected=rejected,
        fingerprint=fingerprint,
        labels=dict(Counter(r["label"] for r in rows)),
    )
    from collections import namedtuple

    return namedtuple("PrepOutputs", ["n_examples", "fingerprint"])(len(rows), fingerprint)


@dsl.component(base_image=BASE_IMAGE)
def load_production(
    registry_dir: str,
) -> NamedTuple("Production", [("prompt_template", str), ("model", str), ("version", str)]):
    """Lit la version en production dans le registre (un fichier JSON sur GCS)."""
    import json
    import os
    from collections import namedtuple

    path = os.path.join(registry_dir, "production.json")
    with open(path, encoding="utf-8") as f:
        prod = json.load(f)
    return namedtuple("Production", ["prompt_template", "model", "version"])(
        prod["prompt_template"], prod["model"], prod["version"]
    )


@dsl.component(base_image=BASE_IMAGE, packages_to_install=[GENAI])
def evaluate_prompt(
    dataset: Input[Dataset],
    prompt_template: str,
    model: str,
    version: str,
    labels: list,
    project: str,
    location: str,
    predictions: Output[Dataset],
    metrics: Output[Metrics],
    confusion: Output[ClassificationMetrics],
):
    """Exécute une version (prompt + modèle) sur le jeu de référence et mesure."""
    import json
    import logging
    import statistics
    import time

    from google import genai
    from google.genai import errors, types

    log = logging.getLogger("evaluate_prompt")
    # project vide : mode clé d'API (exécution locale) ; sinon Vertex AI.
    # Délai maximal par appel : sans lui, une connexion bloquée fige la tâche (et la facture).
    http = types.HttpOptions(timeout=60_000)
    client = (
        genai.Client(vertexai=True, project=project, location=location, http_options=http)
        if project
        else genai.Client(http_options=http)
    )
    config = types.GenerateContentConfig(
        temperature=0,
        response_mime_type="text/x.enum",
        response_schema={"type": "STRING", "enum": labels},
    )

    def classify(text: str):
        # Quota partagé : 6 essais, attente 2, 4, 8, 16, 32 s (environ une minute au total).
        for attempt in range(6):
            try:
                t0 = time.perf_counter()
                r = client.models.generate_content(
                    model=model, contents=prompt_template.format(text=text), config=config
                )
                tokens = r.usage_metadata.total_token_count if r.usage_metadata else 0
                return (r.text or "").strip(), (time.perf_counter() - t0) * 1000, tokens or 0
            except errors.APIError as e:
                if e.code not in (429, 500, 503) or attempt == 5:
                    raise
                wait = 2 ** (attempt + 1)
                log.warning("Gemini %s, nouvel essai dans %ss", e.code, wait)
                time.sleep(wait)

    with open(dataset.path, encoding="utf-8") as f:
        rows = [json.loads(line) for line in f if line.strip()]

    out, latencies, tokens, failures = [], [], 0, 0
    for row in rows:
        try:
            pred, ms, tok = classify(row["text"])
        except Exception as e:  # une erreur ne doit pas annuler toute l'évaluation
            log.error("échec sur %s : %s", row["id"], e)
            pred, ms, tok, failures = "", 0.0, 0, failures + 1
        latencies.append(ms)
        tokens += tok
        out.append({**row, "prediction": pred, "ok": pred == row["label"]})

    accuracy = sum(r["ok"] for r in out) / len(out)
    with open(predictions.path, "w", encoding="utf-8") as f:
        f.writelines(json.dumps(r, ensure_ascii=False) + "\n" for r in out)
    predictions.metadata.update(version=version, model=model)

    metrics.log_metric("accuracy", round(accuracy, 4))
    metrics.log_metric("latency_p50_ms", round(statistics.median(latencies), 1))
    metrics.log_metric("total_tokens", tokens)
    metrics.log_metric("failures", failures)
    metrics.metadata["version"] = version
    metrics.metadata["model"] = model

    matrix = [[0] * len(labels) for _ in labels]
    for r in out:
        if r["prediction"] in labels:
            matrix[labels.index(r["label"])][labels.index(r["prediction"])] += 1
    confusion.log_confusion_matrix(labels, matrix)
    log.info("%s : accuracy=%.3f sur %d exemples", version, accuracy, len(out))


@dsl.component(base_image=BASE_IMAGE)
def compare_versions(
    candidate: Input[Metrics],
    production: Input[Metrics],
    min_accuracy: float,
    min_gain: float,
    max_failures: int,
    verdict: Output[Artifact],
) -> str:
    """Porte de promotion : règles explicites, pas de jugement implicite.

    Retourne « promote » ou « reject » (une seule sortie : kfp.local 2.17 ne sait pas
    évaluer un dsl.If sur un champ de NamedTuple). Le motif part dans l'artefact `verdict`.
    """
    c, p = candidate.metadata, production.metadata
    gain = c["accuracy"] - p["accuracy"]
    if c["failures"] > max_failures:
        decision, reason = "reject", f"{int(c['failures'])} appels en échec (max {max_failures})"
    elif c["accuracy"] < min_accuracy:
        decision, reason = "reject", f"accuracy {c['accuracy']:.2f} < seuil {min_accuracy:.2f}"
    elif gain < min_gain:
        decision, reason = "reject", f"gain {gain:+.2f} < {min_gain:+.2f} exigé"
    else:
        decision, reason = "promote", f"accuracy {p['accuracy']:.2f} -> {c['accuracy']:.2f} ({gain:+.2f})"
    verdict.metadata.update(decision=decision, reason=reason, gain=round(gain, 4),
                            candidate=c.get("version"), production=p.get("version"))
    print(f"Porte : {decision} ({reason})")
    return decision


@dsl.component(base_image=BASE_IMAGE)
def promote(
    registry_dir: str,
    prompt_template: str,
    model: str,
    version: str,
    dataset_fingerprint: str,
    candidate: Input[Metrics],
    verdict: Input[Artifact],
) -> str:
    """Écrit la nouvelle version en production et garde l'ancienne dans l'historique."""
    import datetime as dt
    import json
    import os
    import shutil

    current = os.path.join(registry_dir, "production.json")
    history = os.path.join(registry_dir, "history")
    os.makedirs(history, exist_ok=True)
    if os.path.exists(current):
        with open(current, encoding="utf-8") as f:
            old = json.load(f)
        shutil.copy(current, os.path.join(history, f"{old['version']}.json"))
    record = {
        "version": version,
        "model": model,
        "prompt_template": prompt_template,
        "promoted_at": dt.datetime.now(dt.timezone.utc).isoformat(timespec="seconds"),
        "reason": verdict.metadata["reason"],
        "dataset_fingerprint": dataset_fingerprint,
        "metrics": {k: candidate.metadata[k] for k in ("accuracy", "latency_p50_ms", "total_tokens")},
    }
    tmp = current + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(record, f, ensure_ascii=False, indent=2)
    os.replace(tmp, current)
    return version