import re
from typing import Dict, List

# HIPAA Safe Harbor 18 식별자 중 정규식으로 처리 가능한 주요 패턴
HIPAA_PATTERNS: Dict[str, str] = {
    "SSN": r"\b\d{3}-\d{2}-\d{4}\b",
    "PHONE": r"\b\(?\d{3}\)?[\s.-]?\d{3}[\s.-]?\d{4}\b",
    "EMAIL": r"\b[\w.-]+@[\w.-]+\.\w+\b",
    "MRN": r"\bMRN[:\s]*\d{6,10}\b",
    "DATE_FULL": r"\b\d{4}[-/]\d{1,2}[-/]\d{1,2}\b",
    "DATE_MDY": r"\b\d{1,2}[-/]\d{1,2}[-/]\d{2,4}\b",
    "AGE_90PLUS": r"\b(?:aged?\s+)?(9[0-9]|1[0-2]\d)\s*(?:years?|yrs?|yo)\b",
    "ZIP": r"\b\d{5}(?:-\d{4})?\b",
}


def deidentify(text: str) -> str:
    """HIPAA Safe Harbor 규제 준수 마스킹.
    
    정규식으로 처리 불가한 항목(환자명·지명·기관명)은 SciSpacy NER 병행 필요.
    """
    masked = text
    for label, pattern in HIPAA_PATTERNS.items():
        masked = re.sub(pattern, f"[{label}]", masked, flags=re.IGNORECASE)
    return masked

import json
import anthropic

client = anthropic.Anthropic()

EXTRACTION_SCHEMA = {
    "name": "extract_clinical_fields",
    "description": "임상 노트에서 4개 필드를 정형 JSON으로 추출",
    "input_schema": {
        "type": "object",
        "properties": {
            "symptoms": {
                "type": "array",
                "items": {"type": "string"},
                "description": "환자가 호소한 증상 (예: chest pain, dyspnea)",
            },
            "medications": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "name": {"type": "string"},
                        "dose": {"type": "string"},
                        "frequency": {"type": "string"},
                    },
                    "required": ["name"],
                },
            },
            "labs": {
                "type": "array",
                "items": {
                    "type": "object",
                    "properties": {
                        "test": {"type": "string"},
                        "value": {"type": "string"},
                        "unit": {"type": "string"},
                    },
                    "required": ["test", "value"],
                },
            },
            "diagnoses": {
                "type": "array",
                "items": {"type": "string"},
                "description": "확진·의심 진단명 (ICD-10 코드 또는 자연어)",
            },
        },
        "required": ["symptoms", "medications", "labs", "diagnoses"],
    },
}

SYSTEM_PROMPT = """당신은 임상 자연어 처리 전문가입니다.
주어진 임상 노트에서 4개 필드(symptoms, medications, labs, diagnoses)를 추출합니다.

원칙:
1. 노트에 없는 정보를 추론·창작하지 않습니다 (할루시네이션 금지).
2. 약어는 문맥으로 정확히 해석합니다 (예: MI = myocardial infarction if cardiac context).
3. 마스킹된 토큰([DATE_FULL], [MRN] 등)은 추출 대상이 아닙니다.
4. 불확실한 진단은 diagnoses에서 제외하고 symptoms에만 포함합니다.
"""


def extract_fields(deidentified_note: str) -> dict:
    """Claude API를 호출해 정형 JSON 추출."""
    response = client.messages.create(
        model="claude-sonnet-4-5",
        max_tokens=2048,
        system=SYSTEM_PROMPT,
        tools=[EXTRACTION_SCHEMA],
        tool_choice={"type": "tool", "name": "extract_clinical_fields"},
        messages=[{"role": "user", "content": deidentified_note}],
    )
    for block in response.content:
        if block.type == "tool_use":
            return block.input
    return {}

LAB_PATTERN = re.compile(
    r"(?P<test>[A-Z][a-zA-Z\s]{2,20})\s*[:=]\s*"
    r"(?P<value>\d+\.?\d*)\s*(?P<unit>[a-zA-Z/%]+)?"
)


def regex_fallback_labs(note: str) -> List[dict]:
    """LLM이 labs를 비우면 정규식으로 최소 추출."""
    labs = []
    for m in LAB_PATTERN.finditer(note):
        labs.append({
            "test": m.group("test").strip(),
            "value": m.group("value"),
            "unit": m.group("unit") or "",
        })
    return labs


def merge_with_fallback(llm_output: dict, note: str) -> dict:
    """LLM 결과 + regex fallback."""
    result = dict(llm_output)
    if not result.get("labs"):
        result["labs"] = regex_fallback_labs(note)
    return result

import requests


UMLS_BASE = "https://uts-ws.nlm.nih.gov/rest"


def map_to_umls_cui(term: str, api_key: str) -> str | None:
    """UMLS REST API로 자연어 개념을 CUI로 매핑.
    
    api_key는 UMLS 등록 후 발급받은 값. 무료 등록.
    """
    resp = requests.get(
        f"{UMLS_BASE}/search/current",
        params={"string": term, "apiKey": api_key, "pageSize": 1},
        timeout=10,
    )
    if resp.status_code != 200:
        return None
    results = resp.json().get("result", {}).get("results", [])
    return results[0].get("ui") if results else None

from typing import Set


def f1_score(pred: Set[str], gold: Set[str]) -> float:
    if not pred and not gold:
        return 1.0
    tp = len(pred & gold)
    if tp == 0:
        return 0.0
    precision = tp / len(pred)
    recall = tp / len(gold)
    return 2 * precision * recall / (precision + recall)


def evaluate_extraction(pred_json: dict, gold_json: dict) -> Dict[str, float]:
    """필드별 F1 반환."""
    scores = {}
    for field in ["symptoms", "diagnoses"]:
        pred = {s.lower().strip() for s in pred_json.get(field, [])}
        gold = {s.lower().strip() for s in gold_json.get(field, [])}
        scores[field] = f1_score(pred, gold)
    for field in ["medications", "labs"]:
        pred = {json.dumps(x, sort_keys=True) for x in pred_json.get(field, [])}
        gold = {json.dumps(x, sort_keys=True) for x in gold_json.get(field, [])}
        scores[field] = f1_score(pred, gold)
    return scores

def process_note(raw_note: str, umls_key: str | None = None) -> dict:
    """임상 노트 1건을 5단계 파이프라인으로 처리."""
    deidentified = deidentify(raw_note)
    llm_output = extract_fields(deidentified)
    merged = merge_with_fallback(llm_output, deidentified)
    if umls_key:
        merged["diagnoses_cui"] = [
            map_to_umls_cui(d, umls_key) for d in merged.get("diagnoses", [])
        ]
    return merged