from pydantic import BaseModel
from typing import Optional
from datetime import date


class FHIRCoding(BaseModel):
    system: str
    code: str
    display: str


class FHIRCondition(BaseModel):
    resourceType: str = "Condition"
    id: str
    code: FHIRCoding
    clinicalStatus: str = "active"
    onsetDate: Optional[str] = None


class FHIRObservation(BaseModel):
    resourceType: str = "Observation"
    id: str
    code: FHIRCoding
    valueQuantity: Optional[dict] = None
    valueString: Optional[str] = None
    valueBoolean: Optional[bool] = None
    status: str = "final"


class FHIRMedication(BaseModel):
    resourceType: str = "MedicationStatement"
    id: str
    medication: FHIRCoding
    status: str = "active"


class FHIRPatient(BaseModel):
    resourceType: str = "Patient"
    id: str
    gender: str
    birthDate: str
    conditions: list[FHIRCondition] = []
    observations: list[FHIRObservation] = []
    medications: list[FHIRMedication] = []


def build_patient_profile(fhir_patient: FHIRPatient) -> dict:
    """Convert FHIR R4 patient bundle to normalized matching profile."""
    from datetime import datetime
    birth_year = int(fhir_patient.birthDate[:4])
    age = datetime.now().year - birth_year

    diagnoses = [c.code.code for c in fhir_patient.conditions]
    diagnosis_names = [c.code.display for c in fhir_patient.conditions]
    medications = [m.medication.display for m in fhir_patient.medications]

    biomarkers = {}
    lab_values = {}
    for obs in fhir_patient.observations:
        key = obs.code.display.lower().replace(" ", "_")
        if obs.valueBoolean is not None:
            biomarkers[key] = obs.valueBoolean
        elif obs.valueQuantity:
            lab_values[key] = obs.valueQuantity
        elif obs.valueString:
            biomarkers[key] = obs.valueString

    return {
        "patient_id": fhir_patient.id,
        "age": age,
        "gender": fhir_patient.gender,
        "diagnosis_codes": diagnoses,
        "diagnosis_names": diagnosis_names,
        "medications": medications,
        "biomarkers": biomarkers,
        "lab_values": lab_values,
        "fhir_bundle_ref": f"Patient/{fhir_patient.id}",
    }


# Realistic mock FHIR R4 patients for demo
MOCK_FHIR_PATIENTS: dict[str, FHIRPatient] = {
    "P001": FHIRPatient(
        id="P001", gender="female", birthDate="1979-03-15",
        conditions=[
            FHIRCondition(id="c1", code=FHIRCoding(system="http://snomed.info/sct", code="254837009", display="Breast cancer"), onsetDate="2022-06-01"),
        ],
        observations=[
            FHIRObservation(id="o1", code=FHIRCoding(system="http://loinc.org", code="85319-2", display="HER2"), valueBoolean=True),
            FHIRObservation(id="o2", code=FHIRCoding(system="http://loinc.org", code="2857-1", display="PSA"), valueQuantity={"value": 0.5, "unit": "ng/mL"}),
            FHIRObservation(id="o3", code=FHIRCoding(system="http://loinc.org", code="718-7", display="Hemoglobin"), valueQuantity={"value": 12.5, "unit": "g/dL"}),
        ],
        medications=[
            FHIRMedication(id="m1", medication=FHIRCoding(system="http://www.nlm.nih.gov/research/umls/rxnorm", code="583214", display="Trastuzumab")),
        ],
    ),
    "P002": FHIRPatient(
        id="P002", gender="male", birthDate="1964-08-22",
        conditions=[
            FHIRCondition(id="c2", code=FHIRCoding(system="http://snomed.info/sct", code="399068003", display="Prostate cancer"), onsetDate="2021-11-15"),
        ],
        observations=[
            FHIRObservation(id="o4", code=FHIRCoding(system="http://loinc.org", code="2857-1", display="PSA"), valueQuantity={"value": 8.3, "unit": "ng/mL"}),
            FHIRObservation(id="o5", code=FHIRCoding(system="http://loinc.org", code="85319-2", display="BRCA2"), valueBoolean=True),
        ],
        medications=[
            FHIRMedication(id="m2", medication=FHIRCoding(system="http://www.nlm.nih.gov/research/umls/rxnorm", code="1946819", display="Enzalutamide")),
        ],
    ),
    "P003": FHIRPatient(
        id="P003", gender="female", birthDate="1985-11-30",
        conditions=[
            FHIRCondition(id="c3", code=FHIRCoding(system="http://snomed.info/sct", code="254837009", display="Breast cancer"), onsetDate="2023-02-10"),
            FHIRCondition(id="c4", code=FHIRCoding(system="http://snomed.info/sct", code="44054006", display="Type 2 diabetes"), onsetDate="2019-05-01"),
        ],
        observations=[
            FHIRObservation(id="o6", code=FHIRCoding(system="http://loinc.org", code="85319-2", display="HER2"), valueBoolean=False),
            FHIRObservation(id="o7", code=FHIRCoding(system="http://loinc.org", code="4548-4", display="HbA1c"), valueQuantity={"value": 7.2, "unit": "%"}),
        ],
        medications=[
            FHIRMedication(id="m3", medication=FHIRCoding(system="http://www.nlm.nih.gov/research/umls/rxnorm", code="860975", display="Metformin")),
        ],
    ),
    "P004": FHIRPatient(
        id="P004", gender="male", birthDate="1957-04-07",
        conditions=[
            FHIRCondition(id="c5", code=FHIRCoding(system="http://snomed.info/sct", code="363346000", display="Non-small cell lung cancer"), onsetDate="2022-09-20"),
        ],
        observations=[
            FHIRObservation(id="o8", code=FHIRCoding(system="http://loinc.org", code="81704-9", display="EGFR mutation"), valueString="L858R"),
            FHIRObservation(id="o9", code=FHIRCoding(system="http://loinc.org", code="73977-1", display="PD-L1 expression"), valueQuantity={"value": 60, "unit": "%"}),
        ],
        medications=[
            FHIRMedication(id="m4", medication=FHIRCoding(system="http://www.nlm.nih.gov/research/umls/rxnorm", code="1860492", display="Osimertinib")),
        ],
    ),
    "P005": FHIRPatient(
        id="P005", gender="female", birthDate="1990-07-19",
        conditions=[
            FHIRCondition(id="c6", code=FHIRCoding(system="http://snomed.info/sct", code="93761005", display="Primary malignant neoplasm of colon"), onsetDate="2023-07-01"),
        ],
        observations=[
            FHIRObservation(id="o10", code=FHIRCoding(system="http://loinc.org", code="85077-6", display="MSI status"), valueString="MSI-H"),
            FHIRObservation(id="o11", code=FHIRCoding(system="http://loinc.org", code="85319-2", display="KRAS"), valueBoolean=False),
        ],
        medications=[],
    ),
}


def get_mock_fhir_patient(patient_id: str) -> Optional[FHIRPatient]:
    return MOCK_FHIR_PATIENTS.get(patient_id)


def get_all_patient_ids() -> list[str]:
    return list(MOCK_FHIR_PATIENTS.keys())


def get_patient_profile(patient_id: str) -> Optional[dict]:
    patient = get_mock_fhir_patient(patient_id)
    if patient:
        return build_patient_profile(patient)
    return None