#!/usr/bin/env python3
"""
Extraction des règles du Guide Méthodologique MCO 2026 → Q&A raisonnement.

Le Guide Méthodologique est le document de référence ATIH qui définit :
  - Les règles de codage du DP, DR, DAS
  - Les règles CMA et niveaux de sévérité
  - Les règles de codage des actes CCAM
  - Les cas particuliers (séjours multi-unités, transferts, etc.)

Stratégie :
  - Parser le PDF en sections
  - Extraire les règles de codage
  - Générer des Q&A directes (sans LLM) pour les définitions
  - Générer des Q&A de raisonnement via Claude Opus 4.6 pour les règles
  - Focus sur les règles DP/DAS/CMA les plus applicables

Cible : 500 exemples (183 directs + ~320 LLM)

Source : data/raw/guide_methodo_mco_2026.pdf
Nécessite : ANTHROPIC_API_KEY en variable d'environnement

Usage :
  python scripts/16_parse_guide_metho.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
PDF_PATH = BASE / "data" / "raw" / "guide_methodo_mco_2026.pdf"
OUTPUT = BASE / "data" / "processed" / "guide_metho_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 PMSI. "
    "Tu t'appuies sur le Guide Méthodologique de production des informations "
    "relatives à l'activité MCO (ATIH 2026)."
)

# Patterns pour détecter les règles dans le guide
RULE_PATTERNS = [
    re.compile(r"(?:le\s+)?diagnostic\s+principal\s+(?:est|doit|sera|correspond)", re.IGNORECASE),
    re.compile(r"(?:le\s+)?diagnostic\s+(?:relié|associé)\s+(?:est|doit|sera)", re.IGNORECASE),
    re.compile(r"(?:la\s+)?CMA\s+(?:est|correspond|ne\s+peut)", re.IGNORECASE),
    re.compile(r"(?:le\s+)?DAS\s+(?:est|doit|sera|ne\s+doit\s+pas)", re.IGNORECASE),
    re.compile(r"(?:le\s+)?DP\s+(?:est|doit|sera|ne\s+doit\s+pas)", re.IGNORECASE),
    re.compile(r"(?:le\s+)?DR\s+(?:est|doit|sera)", re.IGNORECASE),
    re.compile(r"(?:on\s+)?code(?:ra)?\s+", re.IGNORECASE),
    re.compile(r"ne\s+(?:pas|doit\s+pas|faut\s+pas)\s+(?:coder|enregistrer|recueillir)", re.IGNORECASE),
    re.compile(r"séjour\s+multi-?uni", re.IGNORECASE),
    re.compile(r"(?:transfert|mutation)\s+", re.IGNORECASE),
    re.compile(r"(?:sévérité|niveau\s+de\s+sévérité|complication)", re.IGNORECASE),
    re.compile(r"(?:ressources?\s+supplémentaires?|consomm)", re.IGNORECASE),
    re.compile(r"(?:groupage|GHM|GHS|CMD)", re.IGNORECASE),
]

GENERATION_PROMPT = """Tu es un formateur DIM. À partir de cette règle du Guide Méthodologique MCO, génère un exercice de raisonnement.

RÈGLE (source : Guide Méthodologique MCO 2026, section « {section} ») :
{rule_text}

Génère un objet JSON avec :
1. "scenario" : un cas clinique réaliste et concis (2-3 phrases) qui illustre l'application de cette règle
2. "reponse" : un objet JSON contenant :
   - "analyse_clinique" : interprétation du cas
   - "regle_pmsi" : la règle du guide méthodologique qui s'applique (citée)
   - "code" : le code CIM-10 ou l'action de codage correcte
   - "confidence" : "high"
   - "justification" : pourquoi cette règle s'applique à ce cas

Réponds UNIQUEMENT avec le JSON, sans texte avant/après."""

# Templates pour les Q&A directes (sans LLM)
DIRECT_QA_TEMPLATES = [
    ("Quelle est la définition du {concept} selon le Guide Méthodologique MCO ?",
     "Selon le Guide Méthodologique MCO 2026 :\n\n{text}"),
    ("Quelles sont les règles de codage pour {concept} selon le Guide Méthodologique ?",
     "Le Guide Méthodologique MCO 2026 précise les règles suivantes pour {concept} :\n\n{text}"),
]


def extract_sections(pdf_path: Path) -> list[tuple[str, str]]:
    """Extraire les sections du Guide Méthodologique."""
    with pdfplumber.open(str(pdf_path)) as pdf:
        full_text = ""
        for page in pdf.pages:
            text = page.extract_text() or ""
            full_text += text + "\n\n"

    lines = []
    for line in full_text.split("\n"):
        stripped = line.strip()
        if re.match(r"^\s*\d{1,3}\s*$", stripped):
            continue
        if re.match(r"^(Guide méthodologique|ATIH|Page\s+\d)", stripped, re.IGNORECASE):
            continue
        if "...." in stripped:
            continue
        lines.append(line)

    clean_text = "\n".join(lines)

    section_pattern = re.compile(
        r"^(\d+(?:\.\d+)*\.?)\s+([A-ZÀÂÉÈÊËÎÏÔÙÛÜÇ].*)",
        re.MULTILINE
    )

    sections = []
    matches = list(section_pattern.finditer(clean_text))

    for i, match in enumerate(matches):
        num = match.group(1).rstrip(".")
        title = match.group(2).strip()
        start = match.end()
        end = matches[i + 1].start() if i + 1 < len(matches) else len(clean_text)
        body = clean_text[start:end].strip()

        if len(body) > 80:
            section_title = f"{num}. {title}"
            sections.append((section_title, body))

    return sections


def extract_rule_paragraphs(sections: list[tuple[str, str]]) -> list[tuple[str, str]]:
    """Extraire les paragraphes contenant des règles de codage."""
    rules = []

    for section_title, body in sections:
        paragraphs = re.split(r"\n\s*\n", body)

        for para in paragraphs:
            para = para.strip()
            if len(para) < 80:
                continue

            matches = sum(1 for pat in RULE_PATTERNS if pat.search(para))
            if matches >= 1:
                text = para[:1500] if len(para) > 1500 else para
                rules.append((section_title, text))

    return rules


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=2048,
                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) -> dict | None:
    """Parse la réponse JSON du LLM."""
    if not response_text:
        return None
    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 "scenario" in data and "reponse" in data:
            return data
    except json.JSONDecodeError:
        pass

    # Fallback
    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 None


def make_chatml(system: str, user: str, assistant: str) -> dict:
    return {
        "messages": [
            {"role": "system", "content": system},
            {"role": "user", "content": user},
            {"role": "assistant", "content": assistant},
        ]
    }


def generate_direct_qa(sections: list[tuple[str, str]]) -> list[dict]:
    """Générer des Q&A directes depuis les sections (sans LLM)."""
    examples = []

    concept_keywords = {
        "diagnostic principal": ["DP", "diagnostic principal"],
        "diagnostic relié": ["DR", "diagnostic relié"],
        "diagnostic associé significatif": ["DAS", "diagnostic associé"],
        "complication ou morbidité associée": ["CMA", "complication", "sévérité"],
        "séjour multi-unités": ["multi-unité", "transfert", "mutation"],
        "groupage GHM": ["GHM", "GHS", "groupage", "CMD"],
        "actes CCAM": ["CCAM", "acte", "procédure"],
    }

    for section_title, body in sections:
        body_lower = body.lower()

        for concept, keywords in concept_keywords.items():
            if any(kw.lower() in body_lower for kw in keywords):
                text = body[:2000] if len(body) > 2000 else body

                tmpl_q, tmpl_a = random.choice(DIRECT_QA_TEMPLATES)
                user = tmpl_q.format(concept=f"{concept} (section {section_title})")
                assistant = tmpl_a.format(concept=concept, text=text)

                examples.append(make_chatml(SYSTEM_PROMPT, user, assistant))
                break

    return examples


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--dry-run", action="store_true", help="Pas d'appel LLM")
    parser.add_argument("--max", type=int, default=500, help="Max exemples à générer via LLM")
    args = parser.parse_args()

    print("=" * 60)
    print("Parsing du Guide Méthodologique MCO 2026")
    print(f"Modèle : {MODEL}")
    print("=" * 60)

    if not PDF_PATH.exists():
        print(f"PDF non trouvé : {PDF_PATH}")
        return

    # 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

    print(f"\nParsing de {PDF_PATH.name} ({PDF_PATH.stat().st_size / 1024 / 1024:.1f} Mo)...")
    sections = extract_sections(PDF_PATH)
    print(f"  Sections extraites : {len(sections)}")

    # Q&A directes (sans LLM)
    print("\nGénération des Q&A directes (sans LLM)...")
    direct_examples = generate_direct_qa(sections)
    print(f"  → {len(direct_examples)} exemples directs")

    # Extraire les paragraphes de règles
    print("\nExtraction des règles de codage...")
    rules = extract_rule_paragraphs(sections)
    print(f"  Règles détectées : {len(rules)}")

    if args.dry_run:
        for i, (section, rule) in enumerate(rules[:20]):
            print(f"  [{i+1}] [{section}] {rule[:100]}...")
        print(f"\n[DRY RUN] {len(rules)} règles à traiter. Relancez sans --dry-run.")
        return

    # Limiter le nombre de règles pour LLM
    if len(rules) > args.max:
        random.shuffle(rules)
        rules = rules[:args.max]

    # Générer les Q&A via Claude
    print(f"\nGénération via Claude ({len(rules)} règles)...")
    llm_examples = []
    n_ok = 0
    n_fail = 0

    for i, (section_title, rule_text) in enumerate(rules):
        prompt = GENERATION_PROMPT.format(section=section_title, rule_text=rule_text)
        response_text = call_claude(client, prompt)
        parsed = parse_llm_response(response_text)

        if parsed and "scenario" in parsed and "reponse" in parsed:
            user_content = (
                f"Cas clinique :\n{parsed['scenario']}\n\n"
                f"Applique les règles du Guide Méthodologique MCO (section « {section_title} »)."
            )
            assistant_content = json.dumps(parsed["reponse"], ensure_ascii=False)
            llm_examples.append(make_chatml(SYSTEM_PROMPT, user_content, assistant_content))
            n_ok += 1
        else:
            n_fail += 1

        if (i + 1) % 50 == 0:
            print(f"  Progression : {i+1}/{len(rules)} (ok={n_ok}, fail={n_fail})")

    print(f"  Résultat LLM : {n_ok} exemples, {n_fail} échecs")

    # Fusionner
    all_examples = direct_examples + llm_examples
    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"  Directs (sans LLM) : {len(direct_examples)}")
    print(f"  LLM (Claude)       : {len(llm_examples)}")
    print(f"Taille : {OUTPUT.stat().st_size / 1024:.0f} Ko")


if __name__ == "__main__":
    main()