"""Core FlatFHIR class and public API."""

from __future__ import annotations

import hashlib
import json
from pathlib import Path
from typing import Any, Optional

import pandas as pd

from flat_fhir.extractors import (
    extract_observation,
    extract_patient,
    extract_medication_request,
    extract_condition,
    extract_allergy,
    extract_procedure,
    extract_imaging,
    extract_immunization,
    extract_encounter,
    extract_diagnostic_report,
)

SHEET_NAMES = [
    "patient", "vitals", "labs", "medications", "conditions",
    "allergies", "procedures", "imaging", "immunizations", "encounters",
    "diagnostic_reports",
]

# Column ordering: primary_key, who, when, then the whats
COMMON_COLUMNS = ["primary_key", "performer", "timestamp", "encounter_id", "source_resource_id", "notes"]


class FlatFHIR:
    """Main class for converting FHIR bundles to flat DataFrames."""

    def __init__(self, source: str | dict | list):
        """Initialize from a file path, a FHIR Bundle dict, or a list of Bundle dicts."""
        if isinstance(source, str):
            source = Path(source)
            with open(source, encoding="utf-8") as f:
                data = json.load(f)
        elif isinstance(source, list):
            # Merge multiple bundles
            data = {"resourceType": "Bundle", "entry": []}
            for item in source:
                if isinstance(item, str):
                    with open(item, encoding="utf-8") as f:
                        item = json.load(f)
                for entry in item.get("entry", []):
                    data["entry"].append(entry)
        else:
            data = source

        self._bundle = data
        self._sheets: dict[str, pd.DataFrame] = {}
        self._process()

    def _process(self):
        """Walk bundle entries and route each resource to its extractor.

        Handles all FHIR R4 bundle types: collection, transaction, batch,
        searchset, document, history. Also handles contained resources.
        """
        rows: dict[str, list[dict]] = {name: [] for name in SHEET_NAMES}

        entries = self._bundle.get("entry", [])
        if not entries and self._bundle.get("resourceType") != "Bundle":
            # Single resource, not a bundle
            entries = [{"resource": self._bundle}]

        for entry in entries:
            # In transaction/batch bundles, resource is nested under "resource"
            # In some entries, the entry IS the resource (e.g., contained)
            resource = entry.get("resource", entry)
            rtype = resource.get("resourceType", "")

            # Skip non-clinical resource types
            if rtype in ("Organization", "Practitioner", "PractitionerRole",
                         "Location", "Claim", "ExplanationOfBenefit",
                         "CareTeam", "CarePlan", "Goal", "Device",
                         "Coverage", "Provenance", "SupplyDelivery"):
                continue

            # Process contained resources if present
            for contained in resource.get("contained", []):
                self._route_resource(contained, rows)

            self._route_resource(resource, rows)

        # Build DataFrames
        for name in SHEET_NAMES:
            if rows[name]:
                df = pd.DataFrame(rows[name])
                # Generate primary keys
                df["primary_key"] = df.apply(
                    lambda r: _make_primary_key(name, r), axis=1
                )
                # Sort by performer (who) then timestamp (when), most recent first
                if name != "patient" and "timestamp" in df.columns:
                    sort_cols = []
                    if "performer" in df.columns:
                        sort_cols.append("performer")
                    sort_cols.append("timestamp")
                    df = df.sort_values(sort_cols, ascending=[True, False] if len(sort_cols) == 2 else [False]).reset_index(drop=True)
                # Reorder columns: primary_key, who, when, then the whats
                if name != "patient":
                    ordered = [c for c in COMMON_COLUMNS if c in df.columns]
                    remaining = sorted([c for c in df.columns if c not in ordered])
                    df = df[ordered + remaining]
                else:
                    # Patient sheet: primary_key first
                    cols = ["primary_key"] + [c for c in df.columns if c != "primary_key"]
                    df = df[cols]
                self._sheets[name] = df
            else:
                self._sheets[name] = pd.DataFrame()

    def _route_resource(self, resource: dict, rows: dict[str, list[dict]]):
        """Route a single FHIR resource to the appropriate extractor."""
        rtype = resource.get("resourceType", "")
        try:
            if rtype == "Patient":
                rows["patient"].append(extract_patient(resource))
            elif rtype == "Observation":
                for sheet, row in extract_observation(resource):
                    rows[sheet].append(row)
            elif rtype == "MedicationRequest":
                rows["medications"].append(extract_medication_request(resource))
            elif rtype == "MedicationStatement":
                rows["medications"].append(extract_medication_request(resource))
            elif rtype == "MedicationAdministration":
                rows["medications"].append(extract_medication_request(resource))
            elif rtype == "Condition":
                rows["conditions"].append(extract_condition(resource))
            elif rtype == "AllergyIntolerance":
                rows["allergies"].append(extract_allergy(resource))
            elif rtype == "Procedure":
                rows["procedures"].append(extract_procedure(resource))
            elif rtype == "ImagingStudy":
                rows["imaging"].append(extract_imaging(resource))
            elif rtype == "Immunization":
                rows["immunizations"].append(extract_immunization(resource))
            elif rtype == "Encounter":
                rows["encounters"].append(extract_encounter(resource))
            elif rtype == "DiagnosticReport":
                rows["diagnostic_reports"].append(extract_diagnostic_report(resource))
        except Exception:
            # Skip malformed resources rather than crashing
            pass

    def __getattr__(self, name: str) -> pd.DataFrame:
        if name in SHEET_NAMES:
            return self._sheets.get(name, pd.DataFrame())
        raise AttributeError(f"'{type(self).__name__}' has no attribute '{name}'")

    def sheets(self) -> list[str]:
        """Return names of non-empty sheets."""
        return [name for name in SHEET_NAMES if not self._sheets[name].empty]

    def to_xlsx(self, output: str, sheets: list[str] | None = None):
        """Write to an XLSX workbook with one sheet per resource type."""
        target_sheets = sheets or self.sheets()
        non_empty = [n for n in target_sheets
                     if not self._sheets.get(n, pd.DataFrame()).empty]
        if not non_empty:
            # Nothing to write — create an empty file with a placeholder sheet
            import openpyxl
            wb = openpyxl.Workbook()
            wb.save(output)
            return
        with pd.ExcelWriter(output, engine="openpyxl") as writer:
            for name in non_empty:
                self._sheets[name].to_excel(writer, sheet_name=name, index=False)

    def to_csv(self, output: str, sheet: str = "vitals"):
        """Write a single sheet to CSV."""
        df = self._sheets.get(sheet, pd.DataFrame())
        df.to_csv(output, index=False)

    def to_parquet(self, output: str):
        """Write all sheets to a single Parquet file (stacked with a 'sheet' column)."""
        frames = []
        for name in self.sheets():
            df = self._sheets[name].copy()
            df.insert(0, "_sheet", name)
            frames.append(df)
        if frames:
            combined = pd.concat(frames, ignore_index=True)
            combined.to_parquet(output, index=False)

    def to_json(self, output: str):
        """Write all sheets to a flat JSON file."""
        data = {}
        for name in self.sheets():
            data[name] = self._sheets[name].to_dict(orient="records")
        with open(output, "w") as f:
            json.dump(data, f, indent=2, default=str)

    def to_fhir(self) -> dict:
        """Round-trip back to a FHIR Bundle. Best-effort reconstruction."""
        from flat_fhir.unflattener import unflatten_sheets
        return unflatten_sheets(self._sheets)

    def to_prompt(self, max_tokens: int = 0) -> str:
        """Generate an LLM-friendly text summary."""
        parts = []

        # Patient header
        if not self._sheets["patient"].empty:
            p = self._sheets["patient"].iloc[0]
            name = f"{p.get('given_name', '')} {p.get('family_name', '')}".strip()
            gender = (p.get("gender", "") or "")[:1].upper()
            dob = p.get("birth_date", "")
            parts.append(f"## Patient: {name}, {gender}, DOB {dob}\n")

        # Active conditions
        if not self._sheets["conditions"].empty:
            active = self._sheets["conditions"]
            active_conds = active[active["clinical_status"].str.lower() == "active"] if "clinical_status" in active.columns else active
            if not active_conds.empty:
                parts.append("## Active Conditions")
                for _, row in active_conds.iterrows():
                    onset = f" (diagnosed {row.get('onset_date', '')})" if row.get('onset_date') else ""
                    parts.append(f"- {row.get('condition_name', 'Unknown')}{onset}")
                parts.append("")

        # Allergies
        if not self._sheets["allergies"].empty:
            parts.append("## Allergies")
            for _, row in self._sheets["allergies"].iterrows():
                rxn = f" (causes {row['reaction']})" if row.get('reaction') else ""
                parts.append(f"- {row.get('substance', 'Unknown')}{rxn}")
            parts.append("")

        # Active medications
        if not self._sheets["medications"].empty:
            meds = self._sheets["medications"]
            active_meds = meds[meds["status"].str.lower() == "active"] if "status" in meds.columns else meds
            if not active_meds.empty:
                parts.append("## Active Medications")
                parts.append("| medication | dose | frequency |")
                parts.append("|---|---|---|")
                for _, row in active_meds.iterrows():
                    dose = f"{row.get('dose', '')} {row.get('dose_unit', '')}".strip()
                    parts.append(f"| {row.get('medication_name', '')} | {dose} | {row.get('frequency', '')} |")
                parts.append("")

        # Vitals
        if not self._sheets["vitals"].empty:
            df = self._sheets["vitals"]
            # Pivot: group by timestamp, merge columns
            vitals_pivot = _pivot_observations(df)
            display_df = vitals_pivot.head(5)
            parts.append(f"## Vitals (last {len(display_df)})")
            parts.append(display_df.to_markdown(index=False))
            parts.append("")

        # Labs
        if not self._sheets["labs"].empty:
            df = self._sheets["labs"]
            labs_pivot = _pivot_observations(df)
            display_df = labs_pivot.head(5)
            parts.append(f"## Labs (last {len(display_df)})")
            parts.append(display_df.to_markdown(index=False))
            parts.append("")

        text = "\n".join(parts)

        if max_tokens > 0:
            # Rough token estimate: 1 token ≈ 4 chars
            max_chars = max_tokens * 4
            if len(text) > max_chars:
                text = text[:max_chars] + "\n... (truncated)"

        return text


def _make_primary_key(sheet: str, row: pd.Series) -> str:
    """Generate a deterministic primary key from sheet + source_resource_id + performer + timestamp."""
    parts = [
        sheet,
        str(row.get("source_resource_id", "")),
        str(row.get("performer", "")),
        str(row.get("timestamp", "")),
    ]
    raw = "|".join(parts)
    return hashlib.sha256(raw.encode()).hexdigest()[:12]


def _pivot_observations(df: pd.DataFrame) -> pd.DataFrame:
    """Merge observation rows that share a timestamp into wide-format rows."""
    if df.empty:
        return df

    value_cols = [c for c in df.columns if c not in
                  ["primary_key", "timestamp", "encounter_id", "performer", "source_resource_id", "notes"]]

    if not value_cols:
        return df

    # Group by timestamp and aggregate non-null values
    grouped = df.groupby("timestamp", sort=False).first().reset_index()
    cols = ["timestamp"] + [c for c in value_cols if c in grouped.columns]
    return grouped[cols]


def flatten(source: str | dict, output: str = "output.xlsx",
            sheets: list[str] | None = None):
    """Flatten a FHIR Bundle to an XLSX file.

    Args:
        source: Path to a FHIR Bundle JSON file or a dict.
        output: Output XLSX file path.
        sheets: Optional list of sheet names to include.
    """
    ff = FlatFHIR(source)
    ff.to_xlsx(output, sheets=sheets)
    return ff


def unflatten(source: str, output: str = "bundle.json"):
    """Convert an XLSX file back to a FHIR Bundle JSON.

    Args:
        source: Path to an XLSX file previously created by flat-fhir.
        output: Output FHIR Bundle JSON file path.
    """
    from flat_fhir.unflattener import unflatten_xlsx
    bundle = unflatten_xlsx(source)
    with open(output, "w") as f:
        json.dump(bundle, f, indent=2, default=str)
    return bundle


def merge(sources: list[str | dict], output: str | None = None) -> FlatFHIR:
    """Merge multiple FHIR Bundles and optionally write to XLSX.

    Args:
        sources: List of file paths or dicts.
        output: Optional output XLSX file path.
    """
    loaded = []
    for src in sources:
        if isinstance(src, str):
            with open(src, encoding="utf-8") as f:
                loaded.append(json.load(f))
        else:
            loaded.append(src)

    ff = FlatFHIR(loaded)
    if output:
        ff.to_xlsx(output)
    return ff
