using System.CommandLine;
using Spectre.Console;

public static class MimicListCommand
{
    public static Command Create()
    {
        var dataDirArg = new Argument<DirectoryInfo>("mimic-dir",
            "Path to directory containing MIMIC-IV CSV files");
        var subjectIdOpt = new Option<int?>("--subject-id", "Filter to a specific patient");
        var stayIdOpt = new Option<int?>("--stay-id", "Filter to a specific ICU stay");

        var command = new Command("mimic-list",
            "List available MIMIC-IV patients and ICU stays")
        {
            dataDirArg, subjectIdOpt, stayIdOpt
        };

        command.SetHandler(context =>
        {
            var dataDir = context.ParseResult.GetValueForArgument(dataDirArg);
            var subjectId = context.ParseResult.GetValueForOption(subjectIdOpt);
            var stayId = context.ParseResult.GetValueForOption(stayIdOpt);

            var loader = new MimicDataLoader(dataDir.FullName);
            var patients = loader.LoadPatients().ToDictionary(p => p.SubjectId);
            var admissions = loader.LoadAdmissions().ToDictionary(a => a.HadmId);
            var stays = loader.LoadIcuStays();

            if (subjectId.HasValue)
                stays = stays.Where(s => s.SubjectId == subjectId.Value).ToList();
            if (stayId.HasValue)
                stays = stays.Where(s => s.StayId == stayId.Value).ToList();

            var table = new Table()
                .Border(TableBorder.Rounded)
                .Title("[bold]MIMIC-IV ICU Stays[/]");

            table.AddColumn("StayId");
            table.AddColumn("SubjectId");
            table.AddColumn("HadmId");
            table.AddColumn("Gender");
            table.AddColumn("Age");
            table.AddColumn("Care Unit");
            table.AddColumn("Admission");
            table.AddColumn("LOS (d)");
            table.AddColumn("Expired");

            foreach (var stay in stays.OrderBy(s => s.SubjectId).ThenBy(s => s.InTime))
            {
                var pt = patients.GetValueOrDefault(stay.SubjectId);
                var adm = admissions.GetValueOrDefault(stay.HadmId);

                table.AddRow(
                    stay.StayId.ToString(),
                    stay.SubjectId.ToString(),
                    stay.HadmId.ToString(),
                    pt?.Gender ?? "?",
                    pt?.AnchorAge.ToString() ?? "?",
                    Markup.Escape(stay.FirstCareUnit),
                    adm?.AdmissionType ?? "?",
                    stay.Los.ToString("F1"),
                    adm?.HospitalExpireFlag == 1 ? "[red]Yes[/]" : "No");
            }

            AnsiConsole.Write(table);

            var patientCount = stays.Select(s => s.SubjectId).Distinct().Count();
            AnsiConsole.MarkupLine(
                $"\nFound [bold]{stays.Count}[/] ICU stays across [bold]{patientCount}[/] patients.");
        });

        return command;
    }
}