#!/usr/bin/env python3
"""
Phase 1D — Génération de données synthétiques via API OpenAI (GPT-4o).

Envoie des métadonnées anonymisées FAISS à un grand modèle pour générer
des exemples de raisonnement DIM complet en format ChatML.

Types d'exemples générés :
  1. Scénario clinique → raisonnement DIM → code CIM-10
  2. Discrimination entre codes proches
  3. Application des règles PMSI (DP/DAS, CMA, exclusions)

Nécessite : OPENAI_API_KEY en variable d'environnement

Usage :
  python scripts/06_generate_synthetic.py [--n 500] [--batch 5] [--model gpt-4o] [--dry-run]
"""

import json
import os
import sys
import time
import random
import argparse
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed

random.seed(42)

BASE = Path(__file__).resolve().parent.parent
T2A = Path("/home/dom/ai/t2a")
OUT = BASE / "data" / "processed"
OUT.mkdir(parents=True, exist_ok=True)

# --- Prompts ---

SYSTEM_PROMPT_SCENARIO = """Tu es un formateur DIM (Département d'Information Médicale) expert en codage PMSI.
Tu génères des scénarios cliniques réalistes et anonymisés pour former des médecins DIM au codage CIM-10.

Pour chaque code CIM-10 fourni, tu dois produire :
1. Un SCÉNARIO CLINIQUE réaliste (3-5 phrases, anonymisé, comme extrait d'un compte-rendu d'hospitalisation)
2. Un RAISONNEMENT DIM structuré montrant la démarche de codage

Le raisonnement doit suivre ces étapes :
- analyse_clinique : ce que le texte clinique révèle
- codes_candidats : 2-3 codes CIM-10 envisageables avec leur libellé
- discrimination : pourquoi le code retenu est le bon (et pas les autres)
- regle_pmsi : règle PMSI applicable (DP/DAS, exclusions, conventions dague/astérisque, etc.)
- code : le code CIM-10 retenu
- confidence : high/medium/low
- justification : synthèse en 1 phrase

IMPORTANT :
- Les scénarios doivent être VARIÉS (âges, sexes, contextes différents)
- Anonymisés (pas de vrais noms/dates)
- Médicalement cohérents
- En français médical professionnel
- Réponse en JSON valide uniquement"""

SYSTEM_PROMPT_DISCRIM = """Tu es un formateur DIM expert en codage PMSI.
Tu crées des exercices de discrimination entre codes CIM-10 proches pour former des médecins DIM.

Pour chaque groupe de codes fourni, génère UN scénario clinique où le choix entre les codes est subtil,
puis montre le raisonnement complet pour arriver au bon code.

IMPORTANT : Réponse en JSON valide uniquement."""

SYSTEM_PROMPT_RULES = """Tu es un formateur DIM expert en règles PMSI.
Tu crées des exercices d'application des règles PMSI (codage DP/DAS, CMA, séjours multi-unités, etc.).

Pour chaque situation fournie, génère un scénario d'hospitalisation et montre comment les règles PMSI
s'appliquent au codage.

IMPORTANT : Réponse en JSON valide uniquement."""


def load_faiss_metadata():
    """Charger les métadonnées FAISS."""
    meta_path = T2A / "data" / "rag_index" / "metadata.json"
    with open(meta_path) as f:
        return json.load(f)


def load_cim10_fhir():
    """Charger les concepts FHIR pour enrichir les prompts."""
    fhir_path = BASE / "data" / "raw" / "smt_cim10_fhir.json"
    if not fhir_path.exists():
        return {}
    with open(fhir_path) as f:
        data = json.load(f)
    by_code = {}
    for c in data.get("concept", []):
        by_code[c["code"]] = c
    return by_code


def load_cocoa_entries():
    """Charger les entrées CoCoA parsées."""
    cocoa_path = OUT / "cocoa_entries_debug.json"
    if not cocoa_path.exists():
        return {}
    with open(cocoa_path) as f:
        return json.load(f)


def clean_extrait(extrait):
    """Nettoyer un extrait FAISS (enlever bruit OCR, numéros de page, etc.)."""
    import re
    # Enlever les numéros de page isolés (sur leur propre ligne ou collés)
    extrait = re.sub(r'\n\s*\d{1,4}\s*\n', '\n', extrait)
    extrait = re.sub(r'^\d{1,4}\s*\n', '', extrait)
    # Enlever les transitions de chapitre
    extrait = re.sub(r'Chapitre\s+[IVX]+\b.*', '', extrait)
    # Enlever les lignes de classification
    extrait = re.sub(r'Classification Internationale.*$', '', extrait, flags=re.MULTILINE)
    # Enlever les lignes vides multiples
    extrait = re.sub(r'\n{2,}', '\n', extrait)
    # Tronquer au premier code d'une AUTRE catégorie
    lines = extrait.split('\n')
    first_code = None
    clean_lines = []
    for line in lines:
        stripped = line.strip()
        # Ligne ne contenant qu'un nombre = bruit
        if re.match(r'^\d{1,4}$', stripped):
            continue
        m = re.match(r'^([A-Z]\d{2}(?:\.\d{1,2})?)[*†]?\s', stripped)
        if m:
            code = m.group(1)
            if first_code is None:
                first_code = code
            elif not code.startswith(first_code[:3]):
                break  # On est passé à un autre groupe de codes
        clean_lines.append(stripped)
    result = '\n'.join(clean_lines).strip()
    # Limiter la longueur
    if len(result) > 400:
        result = result[:400].rsplit('\n', 1)[0]
    return result


def select_codes_for_scenarios(metadata, n=500):
    """Sélectionner les codes CIM-10 les plus intéressants pour la génération."""
    # Filtrer les entrées CIM-10 avec des extraits substantiels
    cim10_entries = [m for m in metadata if m.get("document") == "cim10" and len(m.get("extrait", "")) > 20]

    # Prioriser les codes avec des extraits riches
    cim10_entries.sort(key=lambda x: len(x.get("extrait", "")), reverse=True)

    # Prendre les codes uniques
    seen = set()
    selected = []
    for m in cim10_entries:
        code = m["code"]
        if code not in seen and "." in code:  # Préférer les sous-codes (plus spécifiques)
            seen.add(code)
            selected.append(m)
        if len(selected) >= n:
            break

    # Si pas assez de sous-codes, ajouter des catégories
    if len(selected) < n:
        for m in cim10_entries:
            code = m["code"]
            if code not in seen:
                seen.add(code)
                selected.append(m)
            if len(selected) >= n:
                break

    random.shuffle(selected)
    return selected[:n]


def select_discrimination_groups(metadata, cocoa_entries, n=100):
    """Sélectionner des groupes de codes proches pour la discrimination."""
    # Grouper par catégorie parente (3 premiers caractères)
    by_parent = {}
    for m in metadata:
        if m.get("document") != "cim10":
            continue
        code = m.get("code", "")
        if "." in code:
            parent = code.split(".")[0]
            by_parent.setdefault(parent, []).append(m)

    # Sélectionner les groupes avec 2-6 sous-codes
    groups = []
    for parent, children in by_parent.items():
        if 2 <= len(children) <= 6:
            groups.append({
                "parent": parent,
                "codes": [{"code": c["code"], "extrait": clean_extrait(c["extrait"])[:150]} for c in children]
            })

    random.shuffle(groups)
    return groups[:n]


def build_scenario_prompt(codes_batch, fhir_by_code, cocoa_entries):
    """Construire le prompt pour un batch de codes (scénarios cliniques)."""
    items = []
    for meta in codes_batch:
        code = meta["code"]

        # Source primaire : FHIR (propre)
        fhir = fhir_by_code.get(code, {})
        display = fhir.get("display", "")

        # Source secondaire : CoCoA (riche)
        cocoa = cocoa_entries.get(code, {})
        cocoa_desc = cocoa.get("description", "")
        exclusions = cocoa.get("exclusions", [])[:4]
        synonyms = cocoa.get("synonyms", [])[:5]
        comprend = cocoa.get("comprend", [])[:3]
        severity = cocoa.get("severity")
        clinical = " ".join(cocoa.get("clinical_text", []))[:200]

        # Utiliser la meilleure description disponible
        desc = display or cocoa_desc or clean_extrait(meta["extrait"])[:200]

        item = f"CODE: {code}\nLIBELLÉ: {desc}"
        if synonyms:
            item += f"\nSYNONYMES: {'; '.join(synonyms)}"
        if comprend:
            item += f"\nCOMPREND: {'; '.join(comprend)}"
        if exclusions:
            item += f"\nEXCLUSIONS: {'; '.join(exclusions)}"
        if severity:
            item += f"\nSÉVÉRITÉ CMA: {severity}"
        if clinical:
            item += f"\nDESCRIPTION CLINIQUE: {clinical}"
        items.append(item)

    codes_text = "\n---\n".join(items)

    prompt = f"""Génère un scénario clinique et un raisonnement DIM pour chacun des {len(codes_batch)} codes suivants.

{codes_text}

Réponds en JSON avec cette structure exacte :
{{"scenarios": [
  {{
    "code": "le code CIM-10",
    "scenario_clinique": "Texte du scénario clinique réaliste (3-5 phrases)",
    "raisonnement": {{
      "analyse_clinique": "Analyse des éléments cliniques pertinents",
      "codes_candidats": "2-3 codes envisagés avec libellés",
      "discrimination": "Pourquoi ce code et pas les autres",
      "regle_pmsi": "Règle PMSI applicable",
      "code_retenu": "le code",
      "confidence": "high",
      "justification": "Synthèse en 1 phrase"
    }}
  }},
  ...
]}}

Le tableau "scenarios" DOIT contenir exactement {len(codes_batch)} objets, un par code."""

    return prompt


def build_discrimination_prompt(group, fhir_by_code):
    """Construire le prompt pour un exercice de discrimination."""
    codes_text = "\n".join(
        f"- {c['code']}: {c['extrait']}"
        for c in group["codes"]
    )

    prompt = f"""Voici un groupe de codes CIM-10 de la catégorie {group['parent']} :

{codes_text}

Génère UN scénario clinique réaliste où le choix entre ces codes est subtil et demande une réflexion.
Puis montre le raisonnement DIM complet pour arriver au bon code.

Réponds en JSON :
{{
  "scenario_clinique": "Le scénario (5-8 phrases, réaliste, anonymisé)",
  "codes_en_jeu": ["code1", "code2"],
  "raisonnement": {{
    "analyse_clinique": "...",
    "codes_candidats": "...",
    "discrimination": "Explication détaillée de pourquoi un code est préféré",
    "regle_pmsi": "Règle PMSI applicable",
    "code_retenu": "le code correct",
    "confidence": "high/medium",
    "justification": "..."
  }}
}}

JSON valide uniquement."""

    return prompt


def call_openai(client, model, system_prompt, user_prompt, temperature=0.7):
    """Appeler l'API OpenAI."""
    response = client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=temperature,
        max_tokens=4096,
        response_format={"type": "json_object"},
    )
    content = response.choices[0].message.content
    return json.loads(content)


def _extract_items(result):
    """Extraire la liste d'items depuis la réponse JSON (gère différents formats)."""
    if isinstance(result, list):
        return result
    if isinstance(result, dict):
        # Chercher récursivement un tableau d'items
        for v in result.values():
            if isinstance(v, list) and v and isinstance(v[0], dict):
                return v
        # Si pas de tableau, c'est peut-être un seul item
        if "scenario_clinique" in result or "raisonnement" in result:
            return [result]
        # Dernier recours : aplatir les valeurs dict
        items = []
        for v in result.values():
            if isinstance(v, dict) and ("scenario_clinique" in v or "raisonnement" in v):
                items.append(v)
        if items:
            return items
    return []


def convert_to_chatml(scenario_data):
    """Convertir un résultat de génération en format ChatML."""
    if not isinstance(scenario_data, dict):
        return None

    system_msg = "Tu es un médecin DIM expert en codage CIM-10 pour le PMSI français. Tu codes les diagnostics en suivant une démarche structurée."

    scenario = scenario_data.get("scenario_clinique", "")
    raisonnement = scenario_data.get("raisonnement", {})

    # Si pas de scenario_clinique, chercher dans d'autres clés possibles
    if not scenario:
        scenario = scenario_data.get("scenario", scenario_data.get("texte_clinique", ""))

    # Si le raisonnement est directement dans le dict (pas imbriqué)
    if not raisonnement and "analyse_clinique" in scenario_data:
        raisonnement = {k: v for k, v in scenario_data.items() if k != "scenario_clinique"}

    if not scenario or not raisonnement:
        return None

    user_msg = f"Code ce diagnostic en CIM-10.\n\nTEXTE CLINIQUE : {scenario}"
    assistant_msg = json.dumps(raisonnement, ensure_ascii=False)

    return {
        "messages": [
            {"role": "system", "content": system_msg},
            {"role": "user", "content": user_msg},
            {"role": "assistant", "content": assistant_msg},
        ]
    }


def process_scenario_batch(client, model, batch, fhir_by_code, cocoa_entries, batch_idx):
    """Traiter un batch de codes pour générer des scénarios."""
    prompt = build_scenario_prompt(batch, fhir_by_code, cocoa_entries)

    try:
        result = call_openai(client, model, SYSTEM_PROMPT_SCENARIO, prompt)

        # Le résultat peut être un tableau ou un dict avec une clé contenant le tableau
        items = _extract_items(result)

        examples = []
        for item in items:
            chatml = convert_to_chatml(item)
            if chatml:
                examples.append(chatml)
            else:
                print(f"  [Batch {batch_idx}] Item non converti: {list(item.keys()) if isinstance(item, dict) else type(item)}")

        if len(examples) < len(batch):
            print(f"  [Batch {batch_idx}] {len(examples)}/{len(batch)} exemples récupérés")

        return examples

    except Exception as e:
        print(f"  [Batch {batch_idx}] Erreur: {e}")
        return []


def process_discrimination_batch(client, model, group, fhir_by_code, batch_idx):
    """Traiter un groupe pour générer un exercice de discrimination."""
    prompt = build_discrimination_prompt(group, fhir_by_code)

    try:
        result = call_openai(client, model, SYSTEM_PROMPT_DISCRIM, prompt)
        chatml = convert_to_chatml(result)
        return [chatml] if chatml else []

    except Exception as e:
        print(f"  [Discrim {batch_idx}] Erreur: {e}")
        return []


def main():
    parser = argparse.ArgumentParser(description="Génération de données synthétiques via OpenAI")
    parser.add_argument("--n", type=int, default=500, help="Nombre de scénarios à générer")
    parser.add_argument("--n-discrim", type=int, default=100, help="Nombre d'exercices de discrimination")
    parser.add_argument("--batch", type=int, default=5, help="Codes par batch (scénarios)")
    parser.add_argument("--model", default="gpt-4o", help="Modèle OpenAI")
    parser.add_argument("--workers", type=int, default=3, help="Workers parallèles")
    parser.add_argument("--dry-run", action="store_true", help="Afficher les prompts sans appeler l'API")
    parser.add_argument("--resume", action="store_true", help="Reprendre depuis la dernière exécution")
    args = parser.parse_args()

    # Vérifier la clé API
    api_key = os.environ.get("OPENAI_API_KEY")
    if not api_key and not args.dry_run:
        print("Erreur: OPENAI_API_KEY non définie.")
        print("  export OPENAI_API_KEY='sk-...'")
        sys.exit(1)

    # Charger les données
    print("Chargement des données...")
    metadata = load_faiss_metadata()
    fhir_by_code = load_cim10_fhir()
    cocoa_entries = load_cocoa_entries()
    print(f"  FAISS: {len(metadata)} entrées")
    print(f"  FHIR: {len(fhir_by_code)} concepts")
    print(f"  CoCoA: {len(cocoa_entries)} entrées")

    # Sélectionner les codes
    print(f"\nSélection de {args.n} codes pour les scénarios...")
    selected_codes = select_codes_for_scenarios(metadata, n=args.n)
    print(f"  {len(selected_codes)} codes sélectionnés")

    print(f"Sélection de {args.n_discrim} groupes pour la discrimination...")
    discrim_groups = select_discrimination_groups(metadata, cocoa_entries, n=args.n_discrim)
    print(f"  {len(discrim_groups)} groupes sélectionnés")

    # Découper en batches
    scenario_batches = [
        selected_codes[i:i + args.batch]
        for i in range(0, len(selected_codes), args.batch)
    ]
    print(f"\n{len(scenario_batches)} batches de scénarios ({args.batch} codes/batch)")
    print(f"{len(discrim_groups)} exercices de discrimination")

    # Fichier de sortie (avec reprise possible)
    output_path = OUT / "synthetic_chatml.jsonl"
    existing_count = 0
    if args.resume and output_path.exists():
        with open(output_path) as f:
            existing_count = sum(1 for _ in f)
        print(f"\nReprise: {existing_count} exemples existants")

    if args.dry_run:
        # Mode dry-run : montrer des exemples de prompts
        print("\n=== DRY RUN ===")
        print("\n--- Exemple de prompt scénario ---")
        prompt = build_scenario_prompt(scenario_batches[0], fhir_by_code, cocoa_entries)
        print(prompt[:2000])
        print("\n--- Exemple de prompt discrimination ---")
        if discrim_groups:
            prompt = build_discrimination_prompt(discrim_groups[0], fhir_by_code)
            print(prompt[:1500])
        return

    # Initialiser le client OpenAI
    from openai import OpenAI
    client = OpenAI(api_key=api_key)

    all_examples = []
    total_batches = len(scenario_batches) + len(discrim_groups)
    completed = 0
    errors = 0

    # Ouvrir le fichier en mode append
    mode = "a" if args.resume else "w"
    with open(output_path, mode) as fh:

        # Phase 1 : Scénarios cliniques
        print(f"\n{'='*50}")
        print(f"Phase 1 : Génération des scénarios cliniques...")
        print(f"{'='*50}")

        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = {}
            for i, batch in enumerate(scenario_batches):
                future = executor.submit(
                    process_scenario_batch, client, args.model,
                    batch, fhir_by_code, cocoa_entries, i
                )
                futures[future] = i

            for future in as_completed(futures):
                batch_idx = futures[future]
                try:
                    examples = future.result()
                    for ex in examples:
                        fh.write(json.dumps(ex, ensure_ascii=False) + "\n")
                        all_examples.append(ex)
                    completed += 1
                    if completed % 10 == 0:
                        print(f"  [{completed}/{len(scenario_batches)}] {len(all_examples)} exemples générés...")
                except Exception as e:
                    errors += 1
                    print(f"  [Batch {batch_idx}] Exception: {e}")

        print(f"  Scénarios: {len(all_examples)} exemples")

        # Phase 2 : Discrimination
        print(f"\n{'='*50}")
        print(f"Phase 2 : Génération des exercices de discrimination...")
        print(f"{'='*50}")

        discrim_count = 0
        with ThreadPoolExecutor(max_workers=args.workers) as executor:
            futures = {}
            for i, group in enumerate(discrim_groups):
                future = executor.submit(
                    process_discrimination_batch, client, args.model,
                    group, fhir_by_code, i
                )
                futures[future] = i

            for future in as_completed(futures):
                batch_idx = futures[future]
                try:
                    examples = future.result()
                    for ex in examples:
                        fh.write(json.dumps(ex, ensure_ascii=False) + "\n")
                        all_examples.append(ex)
                        discrim_count += 1
                except Exception as e:
                    errors += 1
                    print(f"  [Discrim {batch_idx}] Exception: {e}")

        print(f"  Discrimination: {discrim_count} exemples")

    # Stats finales
    total = len(all_examples) + existing_count
    print(f"\n{'='*50}")
    print(f"Génération terminée !")
    print(f"  Nouveaux exemples : {len(all_examples)}")
    if existing_count:
        print(f"  Existants (reprise): {existing_count}")
    print(f"  Total : {total}")
    print(f"  Erreurs : {errors}")
    print(f"  Fichier : {output_path}")
    if output_path.exists():
        print(f"  Taille : {output_path.stat().st_size / 1024:.0f} Ko")


if __name__ == "__main__":
    main()