import json
import os
from typing import Dict, Iterable, List, Optional

DEFAULT_TAXONOMY: Dict[str, List[str]] = {
    "oncology": [
        "cancer",
        "tumor",
        "carcinoma",
        "melanoma",
        "leukemia",
        "lymphoma",
        "sarcoma",
        "glioma",
        "myeloma",
    ],
    "neurology": [
        "alzheimer",
        "parkinson",
        "multiple sclerosis",
        "epilepsy",
        "migraine",
        "stroke",
        "neurolog",
    ],
    "immunology": [
        "lupus",
        "rheumatoid",
        "psoriasis",
        "asthma",
        "allergy",
        "autoimmune",
        "immun",
    ],
}


def _load_taxonomy_from_env() -> Optional[Dict[str, List[str]]]:
    raw = os.getenv("SYNTHNEURA_DISEASE_TAXONOMY_JSON")
    if not raw:
        return None
    try:
        data = json.loads(raw)
    except json.JSONDecodeError:
        return None
    if not isinstance(data, dict):
        return None
    cleaned: Dict[str, List[str]] = {}
    for key, value in data.items():
        if not isinstance(key, str) or not isinstance(value, list):
            continue
        terms = [str(item).lower() for item in value if item]
        if terms:
            cleaned[key.lower()] = terms
    return cleaned or None


def load_disease_taxonomy() -> Dict[str, List[str]]:
    return _load_taxonomy_from_env() or DEFAULT_TAXONOMY


def classify_disease(chunks: Iterable[str]) -> Optional[str]:
    taxonomy = load_disease_taxonomy()
    haystack = " ".join([chunk for chunk in chunks if chunk]).lower()
    if not haystack:
        return None
    for category, terms in taxonomy.items():
        if any(term in haystack for term in terms):
            return category
    return "other"
