using Azure;
using Azure.AI.DocumentIntelligence;
using Microsoft.Extensions.Options;

public class AzureDocumentOcrService : IOcrService
{
    private const double DefaultTableConfidence = 0.75;

    private static readonly Dictionary<string, string> ClinicalLabelMap =
        new(StringComparer.OrdinalIgnoreCase)
        {
            ["patient name"] = "patient.fullName",
            ["name"] = "patient.fullName",
            ["full name"] = "patient.fullName",
            ["date of birth"] = "patient.dateOfBirth",
            ["dob"] = "patient.dateOfBirth",
            ["birth date"] = "patient.dateOfBirth",
            ["sex"] = "patient.sex",
            ["gender"] = "patient.sex",
            ["admission date"] = "encounter.admissionDate",
            ["admitted"] = "encounter.admissionDate",
            ["department"] = "encounter.department",
            ["ward"] = "encounter.department",
            ["unit"] = "encounter.department",
            ["room"] = "encounter.roomBed",
            ["bed"] = "encounter.roomBed",
            ["room/bed"] = "encounter.roomBed",
            ["hr"] = "observation.HEART_RATE.value",
            ["heart rate"] = "observation.HEART_RATE.value",
            ["pulse"] = "observation.HEART_RATE.value",
            ["temp"] = "observation.TEMP_C.value",
            ["temperature"] = "observation.TEMP_C.value",
            ["bp sys"] = "observation.BP_SYSTOLIC.value",
            ["systolic"] = "observation.BP_SYSTOLIC.value",
            ["bp dia"] = "observation.BP_DIASTOLIC.value",
            ["diastolic"] = "observation.BP_DIASTOLIC.value",
            ["rr"] = "observation.RESP_RATE.value",
            ["resp rate"] = "observation.RESP_RATE.value",
            ["respiratory rate"] = "observation.RESP_RATE.value",
            ["spo2"] = "observation.SPO2.value",
            ["o2 sat"] = "observation.SPO2.value",
            ["oxygen saturation"] = "observation.SPO2.value",
        };

    private static readonly HashSet<string> TimestampHeaders = new(StringComparer.OrdinalIgnoreCase)
    {
        "time", "date", "datetime", "date/time", "recorded", "recorded at", "timestamp"
    };

    private static readonly HashSet<string> GenericTableHeaders = new(StringComparer.OrdinalIgnoreCase)
    {
        "label", "name", "parameter", "field", "item", "value", "result", "reading"
    };

    private readonly DocumentIntelligenceClient _client;
    private readonly ILogger<AzureDocumentOcrService> _logger;

    public AzureDocumentOcrService(
        IOptions<OcrOptions> options,
        ILogger<AzureDocumentOcrService> logger)
    {
        var opts = options.Value.Azure;
        _client = new DocumentIntelligenceClient(
            new Uri(opts.Endpoint),
            new AzureKeyCredential(opts.ApiKey));
        _logger = logger;
    }

    public async Task<OcrExtractionResult> ExtractAsync(
        Stream documentStream, string contentType)
    {
        var sw = System.Diagnostics.Stopwatch.StartNew();

        var content = BinaryData.FromStream(documentStream);
        var operation = await _client.AnalyzeDocumentAsync(
            WaitUntil.Completed,
            "prebuilt-document",
            content);

        var result = operation.Value;
        var fields = new List<OcrExtractedField>();

        foreach (var kv in result.KeyValuePairs ?? [])
        {
            if (kv.Key?.Content is null || kv.Value?.Content is null) continue;

            var fieldName = MapAzureKeyToFieldName(kv.Key.Content);
            if (fieldName is null) continue;

            fields.Add(new OcrExtractedField(
                fieldName,
                kv.Value.Content,
                kv.Confidence));
        }

        foreach (var table in result.Tables ?? [])
        {
            fields.AddRange(ExtractTableObservations(table));
        }

        var rawText = string.Join("\n", (result.Pages ?? [])
            .SelectMany(p => p.Lines?.Select(l => l.Content) ?? []));

        sw.Stop();
        return new OcrExtractionResult(fields, rawText, (int)sw.ElapsedMilliseconds);
    }

    private static string? MapAzureKeyToFieldName(string key)
    {
        var normalized = key.Trim();
        if (normalized.Length == 0)
            return null;

        return ClinicalLabelMap.GetValueOrDefault(normalized);
    }

    private static List<OcrExtractedField> ExtractTableObservations(DocumentTable table)
    {
        var fields = new List<OcrExtractedField>();
        if (table.Cells is null || table.Cells.Count == 0)
            return fields;

        if (table.ColumnCount == 2)
        {
            var labelValueFields = ExtractLabelValueTable(table);
            if (labelValueFields.Count > 0)
                return labelValueFields;
        }

        return ExtractGridTable(table);
    }

    private static List<OcrExtractedField> ExtractLabelValueTable(DocumentTable table)
    {
        var fields = new List<OcrExtractedField>();
        var dataStartRow = HasGenericHeaderRow(table.Cells) ? 1 : 0;

        for (var row = dataStartRow; row < table.RowCount; row++)
        {
            var key = GetCellContent(table.Cells, row, 0);
            var value = GetCellContent(table.Cells, row, 1);
            if (string.IsNullOrWhiteSpace(key) || string.IsNullOrWhiteSpace(value))
                continue;

            var fieldName = MapAzureKeyToFieldName(key);
            if (fieldName is null)
                continue;

            fields.Add(new OcrExtractedField(fieldName, value.Trim(), DefaultTableConfidence));
        }

        return fields;
    }

    private static List<OcrExtractedField> ExtractGridTable(DocumentTable table)
    {
        var fields = new List<OcrExtractedField>();
        var headerCells = table.Cells
            .Where(c => c.RowIndex == 0)
            .OrderBy(c => c.ColumnIndex)
            .ToList();

        if (headerCells.Count == 0)
            return fields;

        var columnMappings = new Dictionary<int, string?>();
        foreach (var headerCell in headerCells)
        {
            columnMappings[headerCell.ColumnIndex] = MapTableHeader(headerCell.Content);
        }

        var observationCodes = columnMappings.Values
            .Where(v => v is not null and not "recordedAt")
            .Cast<string>()
            .Distinct()
            .ToList();

        if (observationCodes.Count == 0)
            return fields;

        for (var row = 1; row < table.RowCount; row++)
        {
            string? recordedAt = null;

            foreach (var cell in table.Cells.Where(c => c.RowIndex == row))
            {
                if (!columnMappings.TryGetValue(cell.ColumnIndex, out var mapping)
                    || mapping is null
                    || string.IsNullOrWhiteSpace(cell.Content))
                {
                    continue;
                }

                var content = cell.Content.Trim();

                if (mapping == "recordedAt")
                {
                    recordedAt = content;
                    continue;
                }

                fields.Add(new OcrExtractedField(
                    $"observation.{mapping}.value",
                    content,
                    DefaultTableConfidence));
            }

            if (recordedAt is null)
                continue;

            foreach (var code in observationCodes)
            {
                fields.Add(new OcrExtractedField(
                    $"observation.{code}.recordedAt",
                    recordedAt,
                    DefaultTableConfidence));
            }
        }

        return fields;
    }

    private static string? MapTableHeader(string? header)
    {
        if (string.IsNullOrWhiteSpace(header))
            return null;

        var normalized = header.Trim();
        if (TimestampHeaders.Contains(normalized))
            return "recordedAt";

        var fieldName = MapAzureKeyToFieldName(normalized);
        if (fieldName is null)
            return null;

        const string observationPrefix = "observation.";
        const string valueSuffix = ".value";
        if (fieldName.StartsWith(observationPrefix, StringComparison.Ordinal)
            && fieldName.EndsWith(valueSuffix, StringComparison.Ordinal))
        {
            return fieldName[observationPrefix.Length..^valueSuffix.Length];
        }

        return null;
    }

    private static bool HasGenericHeaderRow(IReadOnlyList<DocumentTableCell> cells)
    {
        var headerCells = cells.Where(c => c.RowIndex == 0).ToList();
        if (headerCells.Count != 2)
            return false;

        return headerCells.All(c =>
            GenericTableHeaders.Contains(c.Content?.Trim() ?? string.Empty));
    }

    private static string? GetCellContent(
        IReadOnlyList<DocumentTableCell> cells, int row, int column)
    {
        return cells
            .FirstOrDefault(c => c.RowIndex == row && c.ColumnIndex == column)
            ?.Content;
    }
}