using System.Text.Json;

public record MimicGenerateOptions(
    int? MaxHours = null,
    bool IncludeMedications = true,
    bool IncludeLabs = true);

public class MimicScenarioBuilder
{
    private readonly MimicDataLoader _loader;

    private static readonly string[] ObservationPriority =
    [
        "HEART_RATE", "RESP_RATE", "SYSTOLIC_BP", "DIASTOLIC_BP", "TEMP_C", "SPO2",
        "GCS_EYE", "GCS_VERBAL", "GCS_MOTOR",
        "FIO2_PCT", "PAO2_MMHG",
        "CREATININE_MG_DL", "BILIRUBIN_MG_DL", "PLATELET_K_UL",
        "WBC_K_UL", "POTASSIUM_MEQ_L", "LACTATE_MMOL_L", "GLUCOSE_MG_DL",
        "URINE_OUTPUT_ML_H"
    ];

    public MimicScenarioBuilder(MimicDataLoader loader)
    {
        _loader = loader;
    }

    public (ScenarioFile Scenario, List<string> Warnings) Build(
        MimicIcuStay stay, MimicAdmission admission, MimicPatient patient,
        MimicGenerateOptions options)
    {
        var warnings = new List<string>();
        var scenarioStart = stay.InTime;
        var scenarioEnd = stay.OutTime;

        if (options.MaxHours.HasValue)
        {
            var maxEnd = scenarioStart.AddHours(options.MaxHours.Value);
            if (maxEnd < scenarioEnd)
                scenarioEnd = maxEnd;
        }

        var scenarioPatient = BuildPatient(patient, admission);
        var encounter = BuildEncounter(stay, admission);
        var events = BuildEvents(stay, admission, scenarioStart, scenarioEnd, options, warnings);
        var meta = BuildMeta(stay, admission, patient, scenarioStart, scenarioEnd, events);

        var scenario = new ScenarioFile(meta, scenarioPatient, encounter, events, ExpectedOutcomes: null);
        return (scenario, warnings);
    }

    private static ScenarioPatient BuildPatient(MimicPatient patient, MimicAdmission admission)
    {
        var birthYear = DateTime.UtcNow.Year - patient.AnchorAge;
        var dob = new DateTime(birthYear, 7, 1);

        return new ScenarioPatient(
            FirstName: $"MIMIC-{patient.SubjectId}",
            LastName: $"S{admission.HadmId}",
            DateOfBirth: dob.ToString("yyyy-MM-dd"),
            Gender: patient.Gender == "F" ? "Female" : "Male");
    }

    private static ScenarioEncounter BuildEncounter(MimicIcuStay stay, MimicAdmission admission)
    {
        return new ScenarioEncounter(
            Department: MimicCareUnitMap.ToVigilCareDepartment(stay.FirstCareUnit),
            EncounterType: MimicCareUnitMap.ToVigilCareEncounterType(admission.AdmissionType),
            AttendingPhysician: "MIMIC-Physician",
            RoomBed: $"ICU-{stay.StayId % 100:D2}",
            AdmissionReason: $"MIMIC-IV admission ({admission.AdmissionType}, from {admission.AdmissionLocation ?? "unknown"})");
    }

    private List<ScenarioEvent> BuildEvents(
        MimicIcuStay stay, MimicAdmission admission,
        DateTime scenarioStart, DateTime scenarioEnd,
        MimicGenerateOptions options, List<string> warnings)
    {
        var rawObs = CollectObservations(stay, admission, scenarioStart, scenarioEnd, options);
        var deduplicated = DeduplicateBloodPressure(rawObs);
        var events = new List<ScenarioEvent>();

        foreach (var obs in deduplicated)
        {
            var offsetMinutes = Math.Round((obs.ChartTime - scenarioStart).TotalMinutes);
            if (offsetMinutes < 0) offsetMinutes = 0;

            var data = JsonSerializer.SerializeToElement(new
            {
                code = obs.Code,
                value = obs.Value,
                unit = obs.Unit,
                source = obs.Source
            }, SerializerOptions);

            events.Add(new ScenarioEvent(offsetMinutes, "observation", data, null));
        }

        if (options.IncludeMedications)
        {
            var meds = CollectMedications(admission, scenarioStart, scenarioEnd, warnings);
            events.AddRange(meds);
        }

        events = events.OrderBy(e => e.OffsetMinutes).ToList();
        events = EnforceClusterLimit(events, warnings);
        return events;
    }

    private List<RawObservation> CollectObservations(
        MimicIcuStay stay, MimicAdmission admission,
        DateTime scenarioStart, DateTime scenarioEnd,
        MimicGenerateOptions options)
    {
        var observations = new List<RawObservation>();

        foreach (var ce in _loader.StreamChartEvents(stay.StayId))
        {
            if (ce.ChartTime < scenarioStart || ce.ChartTime > scenarioEnd) continue;
            if (!MimicItemMap.TryMapChartEvent(ce.ItemId, out var mapping)) continue;
            if (ce.ValueNum is null) continue;

            var value = MimicItemMap.ConvertValue(ce.ValueNum.Value, mapping);
            observations.Add(new RawObservation(
                ce.ChartTime, mapping.Code, value, mapping.Unit, mapping.Source, mapping.Priority));
        }

        if (options.IncludeLabs)
        {
            var labCodes = new HashSet<(DateTime time, string code)>(
                observations.Select(o => (o.ChartTime, o.Code)));

            foreach (var le in _loader.StreamLabEvents(admission.HadmId, scenarioStart, scenarioEnd))
            {
                if (!MimicItemMap.TryMapLabEvent(le.ItemId, out var mapping)) continue;
                if (le.ValueNum is null) continue;

                if (labCodes.Contains((le.ChartTime, mapping.Code)))
                    continue;

                observations.Add(new RawObservation(
                    le.ChartTime, mapping.Code, le.ValueNum.Value, mapping.Unit, "Lab", 0));
            }
        }

        return observations.OrderBy(o => o.ChartTime).ToList();
    }

    private static List<RawObservation> DeduplicateBloodPressure(List<RawObservation> observations)
    {
        var bpGroups = observations
            .Where(o => o.Code is "SYSTOLIC_BP" or "DIASTOLIC_BP")
            .GroupBy(o => (Time: RoundToMinute(o.ChartTime), o.Code));

        var removals = new HashSet<RawObservation>();
        foreach (var group in bpGroups)
        {
            var items = group.ToList();
            if (items.Count <= 1) continue;

            var hasPrimary = items.Any(i => i.Priority == 0);
            if (hasPrimary)
            {
                foreach (var fallback in items.Where(i => i.Priority > 0))
                    removals.Add(fallback);
            }
        }

        return removals.Count > 0
            ? observations.Where(o => !removals.Contains(o)).ToList()
            : observations;
    }

    private List<ScenarioEvent> CollectMedications(
        MimicAdmission admission, DateTime scenarioStart, DateTime scenarioEnd,
        List<string> warnings)
    {
        var events = new List<ScenarioEvent>();
        var count = 0;
        var skipped = 0;

        foreach (var rx in _loader.StreamPrescriptions(admission.HadmId, scenarioStart, scenarioEnd))
        {
            var dose = ParseDose(rx.DoseValRx);
            if (dose is null || string.IsNullOrWhiteSpace(rx.DoseUnitRx)
                || string.IsNullOrWhiteSpace(rx.Route))
            {
                skipped++;
                continue;
            }

            var offsetMinutes = Math.Round((rx.StartTime - scenarioStart).TotalMinutes);
            if (offsetMinutes < 0) offsetMinutes = 0;

            var data = JsonSerializer.SerializeToElement(new
            {
                drugName = rx.Drug,
                dose = dose.Value,
                doseUnit = rx.DoseUnitRx,
                route = rx.Route,
                administeredBy = "MIMIC-RN"
            }, SerializerOptions);

            events.Add(new ScenarioEvent(offsetMinutes, "medication", data, null));
            count++;
        }

        if (skipped > 0)
            warnings.Add($"Skipped {skipped} prescriptions with missing dose/unit/route data");

        return events;
    }

    private static List<ScenarioEvent> EnforceClusterLimit(
        List<ScenarioEvent> events, List<string> warnings)
    {
        var result = new List<ScenarioEvent>();
        var clusters = events.GroupBy(e => e.OffsetMinutes).OrderBy(g => g.Key).ToList();
        var spillover = new List<(double offset, ScenarioEvent evt)>();

        foreach (var cluster in clusters)
        {
            var obsInCluster = cluster.Where(e => e.Type == "observation").ToList();
            var otherInCluster = cluster.Where(e => e.Type != "observation").ToList();

            // Add any spillover from previous clusters at this offset
            var spilled = spillover.Where(s => s.offset == cluster.Key).Select(s => s.evt).ToList();
            spillover.RemoveAll(s => s.offset == cluster.Key);
            obsInCluster.AddRange(spilled);

            if (obsInCluster.Count > 10)
            {
                var sorted = obsInCluster
                    .OrderBy(e => GetObservationPriority(e))
                    .ToList();

                var keep = sorted.Take(10).ToList();
                var overflow = sorted.Skip(10).ToList();

                warnings.Add(
                    $"Offset {cluster.Key}: split {obsInCluster.Count} observations " +
                    $"(moved {overflow.Count} to offset {cluster.Key + 1})");

                foreach (var evt in overflow)
                    spillover.Add((cluster.Key + 1,
                        new ScenarioEvent(cluster.Key + 1, evt.Type, evt.Data, evt.Note)));

                obsInCluster = keep;
            }

            result.AddRange(obsInCluster);
            result.AddRange(otherInCluster);
        }

        // Handle any remaining spillover
        foreach (var (offset, evt) in spillover.OrderBy(s => s.offset))
            result.Add(evt);

        return result.OrderBy(e => e.OffsetMinutes).ToList();
    }

    private static int GetObservationPriority(ScenarioEvent evt)
    {
        var code = evt.Data.TryGetProperty("code", out var codeProp)
            ? codeProp.GetString() : null;
        if (code is null) return 999;
        var idx = Array.IndexOf(ObservationPriority, code);
        return idx >= 0 ? idx : 999;
    }

    private static ScenarioMeta BuildMeta(
        MimicIcuStay stay, MimicAdmission admission, MimicPatient patient,
        DateTime scenarioStart, DateTime scenarioEnd,
        List<ScenarioEvent> events)
    {
        var durationMinutes = (int)(scenarioEnd - scenarioStart).TotalMinutes;
        var obsCount = events.Count(e => e.Type == "observation");
        var medCount = events.Count(e => e.Type == "medication");

        var description =
            $"Real de-identified MIMIC-IV data. " +
            $"Subject {patient.SubjectId}, stay {stay.StayId}. " +
            $"{(patient.Gender == "F" ? "Female" : "Male")}, age ~{patient.AnchorAge}. " +
            $"ICU LOS: {stay.Los:F1} days. Care unit: {stay.FirstCareUnit}. " +
            $"{obsCount} observations, {medCount} medications. " +
            (admission.HospitalExpireFlag == 1
                ? "Patient expired during hospitalization."
                : $"Discharged to {admission.DischargeLocation ?? "unknown"}.");

        return new ScenarioMeta(
            Id: $"mimic-s{stay.StayId}",
            Name: $"MIMIC-IV — {stay.FirstCareUnit} ({patient.Gender}, ~{patient.AnchorAge}y)",
            Description: description,
            DurationMinutes: durationMinutes,
            Tags: MimicCareUnitMap.GetTags(stay.FirstCareUnit, admission.HospitalExpireFlag));
    }

    private static decimal? ParseDose(string? doseValRx)
    {
        if (string.IsNullOrWhiteSpace(doseValRx)) return null;

        var val = doseValRx.Trim();
        var dashIdx = val.IndexOf('-');
        if (dashIdx > 0) val = val[..dashIdx];

        return decimal.TryParse(val, System.Globalization.NumberStyles.Any,
            System.Globalization.CultureInfo.InvariantCulture, out var d) && d > 0
            ? d
            : null;
    }

    private static DateTime RoundToMinute(DateTime dt)
        => new(dt.Year, dt.Month, dt.Day, dt.Hour, dt.Minute, 0);

    private static readonly JsonSerializerOptions SerializerOptions = new()
    {
        PropertyNamingPolicy = JsonNamingPolicy.CamelCase
    };

    private record RawObservation(
        DateTime ChartTime, string Code, decimal Value,
        string Unit, string Source, int Priority);
}