import json
from typing import Any, Dict, List

from synthneura.core.schemas import (
    Biomarker,
    ClinicalTrial,
    Endpoint,
    ExtractedTrialData,
)
from synthneura.llm.base import LLMClient

PROMPT_VERSION = "extraction_prompt_v1"

MASTER_PROMPT = (
    "You are an expert clinical-trial intelligence engine operating inside a "
    "regulated, high-stakes pharmaceutical environment.\n\n"
    "Your role is not to summarize text or speculate. You extract, normalize, "
    "and structure clinical meaning with maximum precision, traceability, and "
    "consistency.\n\n"
    "Return schema-compliant JSON only with:\n"
    '{\n  "endpoints": [ ... ],\n  "biomarkers": [ ... ]\n}\n\n'
    "Each endpoint object must include:\n"
    "endpoint_id, raw_text, normalized_name, endpoint_category, endpoint_type, "
    "hierarchy, time_horizon, measurement_context, confidence.\n\n"
    "Each biomarker object must include:\n"
    "biomarker_id, raw_text, normalized_name, variant_or_threshold, "
    "biological_level, context, confidence.\n\n"
    'If information is missing or ambiguous, set confidence to "low" and keep '
    "raw_text."
)

ENDPOINT_CANONICAL_MAP = {
    "overall survival": "Overall Survival (OS)",
    "os": "Overall Survival (OS)",
    "progression-free survival": "Progression-Free Survival (PFS)",
    "progression free survival": "Progression-Free Survival (PFS)",
    "pfs": "Progression-Free Survival (PFS)",
    "objective response rate": "Objective Response Rate (ORR)",
    "overall response rate": "Objective Response Rate (ORR)",
    "orr": "Objective Response Rate (ORR)",
    "disease-free survival": "Disease-Free Survival (DFS)",
    "dfs": "Disease-Free Survival (DFS)",
    "time to progression": "Time to Progression (TTP)",
    "ttp": "Time to Progression (TTP)",
    "adverse events": "Safety / Adverse Events",
    "safety": "Safety / Adverse Events",
    "quality of life": "Quality of Life",
}

BIOMARKER_CANONICAL_MAP = {
    "egfr": "EGFR",
    "erbb1": "EGFR",
    "pd-l1": "PD-L1",
    "pdl1": "PD-L1",
    "cd274": "PD-L1",
    "kras": "KRAS",
    "k-ras": "KRAS",
    "braf": "BRAF",
    "alk": "ALK",
    "her2": "HER2",
    "erbb2": "HER2",
    "ros1": "ROS1",
    "ntrk": "NTRK",
}


def _canonicalize(value: str, mapping: Dict[str, str]) -> str:
    key = value.strip().lower()
    return mapping.get(key, value.strip())


def _confidence_to_score(confidence: str) -> float:
    value = confidence.lower()
    if value == "high":
        return 1.0
    if value == "medium":
        return 0.66
    return 0.33


def _endpoint_category(normalized: str) -> str:
    lowered = normalized.lower()
    if any(token in lowered for token in ("survival", "response", "progression")):
        return "efficacy"
    if "quality of life" in lowered:
        return "quality_of_life"
    if "adverse" in lowered or "safety" in lowered:
        return "safety"
    return "other"


def _endpoint_type(normalized: str) -> str:
    lowered = normalized.lower()
    if any(token in lowered for token in ("survival", "progression", "time")):
        return "time_to_event"
    if "rate" in lowered or "response" in lowered:
        return "binary"
    return "unknown"


def _normalize_endpoint(raw: Any, trial: ClinicalTrial, index: int) -> Endpoint:
    if isinstance(raw, dict):
        raw_text = str(raw.get("raw_text") or raw.get("normalized_name") or "")
        normalized = _canonicalize(
            str(raw.get("normalized_name") or raw_text), ENDPOINT_CANONICAL_MAP
        )
        return Endpoint(
            endpoint_id=str(raw.get("endpoint_id") or f"{trial.nct_id}-ep-{index}"),
            raw_text=raw_text,
            normalized_name=normalized,
            endpoint_category=str(
                raw.get("endpoint_category") or _endpoint_category(normalized)
            ),
            endpoint_type=str(raw.get("endpoint_type") or _endpoint_type(normalized)),
            hierarchy=str(raw.get("hierarchy") or "unspecified"),
            time_horizon=raw.get("time_horizon"),
            measurement_context=raw.get("measurement_context"),
            confidence=str(raw.get("confidence") or "low"),
        )
    raw_text = str(raw)
    normalized = _canonicalize(raw_text, ENDPOINT_CANONICAL_MAP)
    return Endpoint(
        endpoint_id=f"{trial.nct_id}-ep-{index}",
        raw_text=raw_text,
        normalized_name=normalized,
        endpoint_category=_endpoint_category(normalized),
        endpoint_type=_endpoint_type(normalized),
        hierarchy="unspecified",
        time_horizon=None,
        measurement_context=None,
        confidence="low",
    )


def _normalize_biomarker(raw: Any, trial: ClinicalTrial, index: int) -> Biomarker:
    if isinstance(raw, dict):
        raw_text = str(raw.get("raw_text") or raw.get("normalized_name") or "")
        normalized = _canonicalize(
            str(raw.get("normalized_name") or raw_text), BIOMARKER_CANONICAL_MAP
        )
        return Biomarker(
            biomarker_id=str(raw.get("biomarker_id") or f"{trial.nct_id}-bm-{index}"),
            raw_text=raw_text,
            normalized_name=normalized,
            variant_or_threshold=raw.get("variant_or_threshold"),
            biological_level=str(raw.get("biological_level") or "gene"),
            context=str(raw.get("context") or "outcome"),
            confidence=str(raw.get("confidence") or "low"),
        )
    raw_text = str(raw)
    normalized = _canonicalize(raw_text, BIOMARKER_CANONICAL_MAP)
    return Biomarker(
        biomarker_id=f"{trial.nct_id}-bm-{index}",
        raw_text=raw_text,
        normalized_name=normalized,
        variant_or_threshold=None,
        biological_level="gene",
        context="outcome",
        confidence="low",
    )


def _build_prompt(trial: ClinicalTrial) -> str:
    """
    Build a concise prompt from normalized trial fields.

    The prompt is optimized for endpoint and biomarker extraction.
    """

    parts: List[str] = []
    if trial.title:
        parts.append(f"Title: {trial.title}")
    if trial.conditions:
        parts.append(f"Conditions: {', '.join(trial.conditions)}")
    if trial.interventions:
        parts.append(f"Interventions: {', '.join(trial.interventions)}")
    if trial.outcomes:
        parts.append(f"Outcomes: {', '.join(trial.outcomes)}")
    content = "\n".join(parts)
    return f"{MASTER_PROMPT}\n\n{content}"


def extract_trial_facts(
    trial: ClinicalTrial,
    client: LLMClient,
) -> ExtractedTrialData:
    """
    Run LLM extraction for a single trial and return structured facts.
    """

    prompt = _build_prompt(trial)
    raw = client.extract_endpoints_biomarkers(prompt)
    endpoints: List[Endpoint] = []
    biomarkers: List[Biomarker] = []
    validation_errors: List[str] = []
    validation_status = "valid"
    payload: Dict[str, Any] = {}
    try:
        payload = json.loads(raw.raw_json)
    except json.JSONDecodeError as exc:
        validation_status = "invalid"
        validation_errors.append(f"json_decode_error: {exc}")

    if validation_status == "valid":
        raw_endpoints = payload.get("endpoints") if isinstance(payload, dict) else None
        raw_biomarkers = (
            payload.get("biomarkers") if isinstance(payload, dict) else None
        )
        if not isinstance(raw_endpoints, list):
            validation_status = "invalid"
            validation_errors.append("endpoints must be a list")
            raw_endpoints = []
        if not isinstance(raw_biomarkers, list):
            validation_status = "invalid"
            validation_errors.append("biomarkers must be a list")
            raw_biomarkers = []
        for idx, item in enumerate(raw_endpoints, start=1):
            try:
                endpoint = _normalize_endpoint(item, trial, idx)
                if not endpoint.raw_text:
                    raise ValueError("missing raw_text")
                endpoints.append(endpoint)
            except Exception as exc:
                validation_status = "invalid"
                validation_errors.append(f"endpoint[{idx}]: {exc}")
        for idx, item in enumerate(raw_biomarkers, start=1):
            try:
                biomarker = _normalize_biomarker(item, trial, idx)
                if not biomarker.raw_text:
                    raise ValueError("missing raw_text")
                biomarkers.append(biomarker)
            except Exception as exc:
                validation_status = "invalid"
                validation_errors.append(f"biomarker[{idx}]: {exc}")

    if endpoints or biomarkers:
        scores = [_confidence_to_score(item.confidence) for item in endpoints] + [
            _confidence_to_score(item.confidence) for item in biomarkers
        ]
        confidence = round(sum(scores) / max(len(scores), 1), 2)
    else:
        confidence = None

    return ExtractedTrialData(
        endpoints=endpoints,
        biomarkers=biomarkers,
        prompt_version=PROMPT_VERSION,
        raw_response=raw.raw_json,
        validation_status=validation_status,
        validation_errors=validation_errors,
        confidence=confidence,
        model=raw.model,
        provider=raw.provider,
    )
