#!/usr/bin/env python3
"""
Parsing agressif des 10 fascicules ATIH → Q&A de raisonnement DIM.

Stratégie :
  - Découpe chaque fascicule en paragraphes (pas sections)
  - Extrait les règles de codage via regex
  - Pour chaque règle extraite, génère 3 Q&A raisonnement via Claude Opus 4.6
  - Question = scénario clinique appliquant la règle
  - Réponse = JSON structuré {analyse_clinique, regle_pmsi, code, justification}

Cible : ~450 exemples (151 règles × 3 exercices)

Sources : data/raw/referentiels/ (10 fascicules PDF, via lien t2a)
Nécessite : ANTHROPIC_API_KEY en variable d'environnement

Usage :
  python scripts/13_generate_fascicule_reasoning.py [--dry-run] [--max N]
"""

import argparse
import json
import os
import random
import re
import sys
import time
from pathlib import Path

import pdfplumber

random.seed(42)

BASE = Path(__file__).resolve().parent.parent
T2A = BASE.parent / "t2a"
REF_DIR = T2A / "data" / "referentiels"
OUTPUT = BASE / "data" / "processed" / "fascicule_reasoning_chatml.jsonl"
OUTPUT.parent.mkdir(parents=True, exist_ok=True)

MODEL = "claude-opus-4-6"

SYSTEM_PROMPT = (
    "Tu es un médecin DIM expert en codage CIM-10 pour le PMSI français. "
    "Tu t'appuies sur les référentiels ATIH officiels."
)

# Fascicules à parser
FASCICULES = {
    "Fascicule_01_Generalites": "Généralités du codage PMSI",
    "Fascicule_02_Maladies_digestives": "Maladies de l'appareil digestif",
    "Fascicule_03_Tumeurs": "Tumeurs",
    "Fascicule_04_Metabolisme": "Métabolisme",
    "Fascicule_05_Gyneco_Obstetrique": "Gynécologie et Obstétrique",
    "Fascicule_06_Neonatalogie": "Néonatalogie",
    "Fascicule_07_Evolutions_2010": "Évolutions 2010",
    "Fascicule_08_Maladies_infectieuses": "Maladies infectieuses",
    "Fascicule_09_AVC": "Accidents vasculaires cérébraux",
    "Fascicule_10_SCA_Coronariens": "Syndromes coronariens aigus",
}

# Regex pour détecter les règles de codage dans le texte
RULE_PATTERNS = [
    re.compile(r"(?:on\s+code(?:ra)?|se\s+code|coder?|est\s+cod[ée](?:e)?)\s+(?:avec\s+|par\s+|en\s+)?([A-Z]\d{2}(?:\.\d{1,2})?)", re.IGNORECASE),
    re.compile(r"ne\s+(?:pas|doit\s+pas|faut\s+pas)\s+coder", re.IGNORECASE),
    re.compile(r"à\s+l['\u2019]exclusion\s+de", re.IGNORECASE),
    re.compile(r"(?:comprend|inclut|inclus)\s+", re.IGNORECASE),
    re.compile(r"(?:dans\s+ce\s+cas|en\s+cas\s+de|si\s+le\s+patient|lorsque)", re.IGNORECASE),
    re.compile(r"(?:le\s+DP|diagnostic\s+principal)\s+(?:est|sera|doit\s+être)", re.IGNORECASE),
    re.compile(r"(?:le\s+DAS|diagnostic\s+associé)\s+(?:est|sera|doit\s+être)", re.IGNORECASE),
    re.compile(r"(?:CMA|sévérité|niveau\s+\d)", re.IGNORECASE),
    re.compile(r"(?:séjour|durée|ressources?\s+supplémentaires?)", re.IGNORECASE),
]

GENERATION_PROMPT = """Tu es un formateur DIM. À partir de cet extrait du fascicule ATIH, génère EXACTEMENT 3 exercices de raisonnement distincts.

EXTRAIT (source : fascicule ATIH « {topic} ») :
{rule_text}

Génère un tableau JSON de 3 objets, chacun avec :
- "scenario" : un cas clinique réaliste et concis (2-3 phrases), DIFFÉRENT des autres
- "reponse" : un objet contenant :
  - "analyse_clinique" : interprétation du cas
  - "regle_pmsi" : la règle du fascicule qui s'applique
  - "code" : le code CIM-10 correct (ou null si ne pas coder)
  - "confidence" : "high"
  - "justification" : pourquoi cette règle s'applique

Réponds UNIQUEMENT avec le tableau JSON (pas d'objet wrapper), sans texte avant/après.
Exemple de format : [{{"scenario": "...", "reponse": {{...}}}}, ...]"""


def extract_paragraphs(pdf_path: Path) -> list[str]:
    """Extraire les paragraphes d'un fascicule PDF.

    Stratégie : pdfplumber ne produit pas de lignes vides entre paragraphes.
    On découpe par page, puis par titres/sections détectées via heuristiques.
    """
    pages_text = []
    with pdfplumber.open(str(pdf_path)) as pdf:
        for page in pdf.pages:
            text = page.extract_text() or ""
            if text.strip():
                pages_text.append(text)

    # Concaténer avec séparateur de page
    full_text = "\n\n".join(pages_text)

    # Nettoyer les lignes non pertinentes
    lines = []
    for line in full_text.split("\n"):
        stripped = line.strip()
        if "...." in stripped or "TABLE DES MATIERES" in stripped:
            continue
        if re.match(r"^\s*\d{1,3}\s*$", stripped):
            continue
        if re.match(r"^(Créé le|Mis à jour|Modifié le|FASCICULE DE CODAGE)", stripped):
            continue
        lines.append(line)

    clean_text = "\n".join(lines)

    # Découper en sections par les titres détectés
    title_re = re.compile(
        r"^(?:"
        r"[IVX]+\.\d*\.?\s+"
        r"|[IVX]+\s+[-–]\s+"
        r"|\d+\.\d*\.?\s+[A-ZÀÂÉÈÊËÎÏ]"
        r"|[A-ZÀÂÉÈÊËÎÏÔÙÛÜ][A-ZÀÂÉÈÊËÎÏÔÙÛÜ\s]{3,}$"
        r")",
        re.MULTILINE
    )

    paragraphs = []

    matches = list(title_re.finditer(clean_text))
    if matches:
        for i, match in enumerate(matches):
            start = match.start()
            end = matches[i + 1].start() if i + 1 < len(matches) else len(clean_text)
            section = clean_text[start:end].strip()
            if len(section) > 100:
                paragraphs.append(section)

    # Fallback : découper par blocs de ~800 caractères
    if len(paragraphs) < 5:
        paragraphs = []
        block_size = 800
        current_block = []
        current_len = 0
        for line in clean_text.split("\n"):
            current_block.append(line)
            current_len += len(line) + 1
            if current_len >= block_size:
                para = "\n".join(current_block).strip()
                if len(para) > 100:
                    paragraphs.append(para)
                current_block = []
                current_len = 0
        if current_block:
            para = "\n".join(current_block).strip()
            if len(para) > 100:
                paragraphs.append(para)

    return paragraphs


def extract_rules(paragraphs: list[str]) -> list[tuple[str, int]]:
    """Extraire les paragraphes classés par pertinence (score de règle).

    Returns: [(text, score)] trié par score décroissant.
    """
    scored = []
    for para in paragraphs:
        score = sum(1 for pat in RULE_PATTERNS if pat.search(para))
        codes = re.findall(r"[A-Z]\d{2}(?:\.\d{1,2})?", para)
        if codes:
            score += 1
        if score >= 1:
            text = para[:1500] if len(para) > 1500 else para
            scored.append((text, score))
    scored.sort(key=lambda x: -x[1])
    return scored


def call_claude(client, prompt: str, max_retries: int = 2) -> str | None:
    """Appel Claude Opus 4.6 via API Anthropic avec retry."""
    for attempt in range(max_retries + 1):
        try:
            response = client.messages.create(
                model=MODEL,
                max_tokens=4096,
                temperature=0.7,
                messages=[{"role": "user", "content": prompt}],
            )
            return response.content[0].text
        except Exception as e:
            if attempt < max_retries:
                wait = 2 ** (attempt + 1)
                print(f"    Retry in {wait}s: {e}")
                time.sleep(wait)
            else:
                print(f"  Claude error: {e}")
                return None


def parse_llm_response(response_text: str) -> list[dict]:
    """Parse la réponse JSON du LLM (tableau de 3 exercices ou objet unique)."""
    if not response_text:
        return []
    text = response_text.strip()
    if "```json" in text:
        text = text.split("```json", 1)[1].split("```", 1)[0].strip()
    elif "```" in text:
        text = text.split("```", 1)[1].split("```", 1)[0].strip()

    try:
        data = json.loads(text)
        if isinstance(data, list):
            return [d for d in data if isinstance(d, dict) and "scenario" in d and "reponse" in d]
        if isinstance(data, dict) and "scenario" in data and "reponse" in data:
            return [data]
    except json.JSONDecodeError:
        pass

    # Fallback : chercher un tableau [...]
    bracket_start = text.find("[")
    if bracket_start >= 0:
        depth = 0
        for i in range(bracket_start, len(text)):
            if text[i] == "[":
                depth += 1
            elif text[i] == "]":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(text[bracket_start:i+1])
                        if isinstance(data, list):
                            return [d for d in data if isinstance(d, dict) and "scenario" in d]
                    except json.JSONDecodeError:
                        break

    # Fallback : objet unique
    brace_start = text.find("{")
    if brace_start >= 0:
        depth = 0
        for i in range(brace_start, len(text)):
            if text[i] == "{":
                depth += 1
            elif text[i] == "}":
                depth -= 1
                if depth == 0:
                    try:
                        data = json.loads(text[brace_start:i+1])
                        if "scenario" in data:
                            return [data]
                    except json.JSONDecodeError:
                        break
    return []


def make_chatml(scenario: str, response: dict, topic: str) -> dict:
    """Créer un exemple ChatML depuis le scénario + réponse structurée."""
    user_content = (
        f"Cas clinique :\n{scenario}\n\n"
        f"Code ce cas en CIM-10 selon les règles du fascicule « {topic} ».\n\n"
        "Réponds avec un JSON structuré contenant : analyse_clinique, regle_pmsi, code, confidence, justification."
    )
    assistant_content = json.dumps(response, ensure_ascii=False)

    return {
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_content},
            {"role": "assistant", "content": assistant_content},
        ]
    }


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true", help="Pas d'appel LLM, affiche les règles extraites")
    parser.add_argument("--max", type=int, default=0, help="Max règles par fascicule (0=illimité)")
    args = parser.parse_args()

    print("=" * 60)
    print("Génération de Q&A raisonnement depuis les fascicules ATIH")
    print(f"Modèle : {MODEL}")
    print("=" * 60)

    # Vérifier la clé API
    if not args.dry_run:
        api_key = os.environ.get("ANTHROPIC_API_KEY")
        if not api_key:
            print("Erreur: ANTHROPIC_API_KEY non définie.")
            print("  export ANTHROPIC_API_KEY='sk-ant-...'")
            sys.exit(1)
        import anthropic
        client = anthropic.Anthropic(api_key=api_key)
    else:
        client = None

    all_examples = []
    total_rules = 0

    for filename_part, topic in FASCICULES.items():
        pdf_files = list(REF_DIR.glob(f"*{filename_part}*"))
        pdf_files = [f for f in pdf_files if "redacted" not in f.name.lower() and "pseudonymise" not in str(f)]
        if not pdf_files:
            print(f"\n{topic} : PDF non trouvé, skip")
            continue

        pdf_path = pdf_files[0]
        print(f"\n{'─'*40}")
        print(f"{topic} ({pdf_path.name})")

        paragraphs = extract_paragraphs(pdf_path)
        print(f"  Paragraphes extraits : {len(paragraphs)}")

        scored_rules = extract_rules(paragraphs)
        print(f"  Règles/paragraphes pertinents : {len(scored_rules)}")

        if args.max > 0:
            scored_rules = scored_rules[:args.max]

        total_rules += len(scored_rules)

        if args.dry_run:
            for i, (rule, score) in enumerate(scored_rules[:5]):
                print(f"  [{i+1}] (score={score}) {rule[:120]}...")
            continue

        # Générer les Q&A via Claude (3 exercices par règle)
        n_ok = 0
        n_fail = 0
        for i, (rule_text, score) in enumerate(scored_rules):
            prompt = GENERATION_PROMPT.format(topic=topic, rule_text=rule_text)
            response_text = call_claude(client, prompt)
            exercises = parse_llm_response(response_text)

            for ex in exercises:
                if "scenario" in ex and "reponse" in ex:
                    example = make_chatml(ex["scenario"], ex["reponse"], topic)
                    all_examples.append(example)
                    n_ok += 1

            if not exercises:
                n_fail += 1

            if (i + 1) % 10 == 0:
                print(f"  Progression : {i+1}/{len(scored_rules)} (exemples={n_ok}, échecs={n_fail})")

        print(f"  Résultat : {n_ok} exemples générés, {n_fail} échecs")

    if args.dry_run:
        print(f"\n[DRY RUN] Total règles détectées : {total_rules}")
        print("Relancez sans --dry-run pour générer les exemples avec Claude.")
        return

    # Mélanger et sauvegarder
    random.shuffle(all_examples)

    with open(OUTPUT, "w") as f:
        for ex in all_examples:
            f.write(json.dumps(ex, ensure_ascii=False) + "\n")

    print(f"\n{'='*60}")
    print(f"Total : {len(all_examples)} exemples → {OUTPUT}")
    print(f"Taille : {OUTPUT.stat().st_size / 1024:.0f} Ko")
    print(f"Règles sources : {total_rules}")


if __name__ == "__main__":
    main()