using System.Text.Json;
using Microsoft.Extensions.Options;

namespace MeetingAssistant.Transcription;

public sealed class FunAsrTranscriptFinalizer
{
    private const string JsonStart = "__MEETING_ASSISTANT_DIARIZATION_JSON_START__";
    private const string JsonEnd = "__MEETING_ASSISTANT_DIARIZATION_JSON_END__";

    private readonly ICommandRunner commandRunner;
    private readonly MeetingAssistantOptions options;
    private readonly ILogger<FunAsrTranscriptFinalizer> logger;

    public FunAsrTranscriptFinalizer(
        ICommandRunner commandRunner,
        IOptions<MeetingAssistantOptions> options,
        ILogger<FunAsrTranscriptFinalizer> logger)
    {
        this.commandRunner = commandRunner;
        this.options = options.Value;
        this.logger = logger;
    }

    public async Task<IReadOnlyList<TranscriptionSegment>> FinalizeAsync(
        string audioPath,
        IReadOnlyList<TranscriptionSegment> liveSegments,
        CancellationToken cancellationToken)
    {
        var diarization = options.FunAsr.Diarization;
        if (!diarization.Enabled)
        {
            return [];
        }

        var fullAudioPath = Path.GetFullPath(Environment.ExpandEnvironmentVariables(audioPath));
        if (!File.Exists(fullAudioPath))
        {
            throw new FileNotFoundException($"Recorded audio was not found at '{fullAudioPath}'.", fullAudioPath);
        }

        var result = await RunDiarizationAsync(fullAudioPath, cancellationToken);
        if (result.ExitCode != 0)
        {
            throw new InvalidOperationException(
                $"FunASR diarization failed with exit code {result.ExitCode}: {result.StandardError}");
        }

        var json = ExtractJson(result.StandardOutput);
        var segments = ParseSegments(json);
        logger.LogInformation("FunASR final diarization produced {SegmentCount} sentence segments", segments.Count);
        return segments;
    }

    private async Task<CommandResult> RunDiarizationAsync(
        string fullAudioPath,
        CancellationToken cancellationToken)
    {
        var backend = options.FunAsr.Backend;
        var modelsFolder = VaultPath.Resolve(backend.ModelsFolder);
        Directory.CreateDirectory(modelsFolder);
        using var timeoutSource = options.FunAsr.Diarization.CommandTimeout > TimeSpan.Zero
            ? new CancellationTokenSource(options.FunAsr.Diarization.CommandTimeout)
            : null;
        using var linkedSource = timeoutSource is null
            ? null
            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);

        try
        {
            return await commandRunner.RunAsync(
                backend.DockerCommand,
                BuildDockerArguments(fullAudioPath, modelsFolder),
                linkedSource?.Token ?? cancellationToken);
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
        {
            throw new TimeoutException(
                $"FunASR diarization timed out after {options.FunAsr.Diarization.CommandTimeout}.");
        }
    }

    private string[] BuildDockerArguments(string fullAudioPath, string modelsFolder)
    {
        var backend = options.FunAsr.Backend;
        return
        [
            "run",
            "--rm",
            "-v",
            $"{fullAudioPath}:/workspace/input.wav:ro",
            "-v",
            $"{modelsFolder}:/workspace/models",
            backend.Image,
            "bash",
            "-lc",
            BuildPythonCommand()
        ];
    }

    private string BuildPythonCommand()
    {
        var diarization = options.FunAsr.Diarization;
        var presetSpeakerCount = diarization.PresetSpeakerCount is int value
            ? $", preset_spk_num={value}"
            : "";
        return
            "cat > /tmp/meeting_assistant_diarize.py <<'PY'\n"
            + "import json\n"
            + "from funasr import AutoModel\n"
            + "model = AutoModel(\n"
            + $"    model={JsonSerializer.Serialize(diarization.AsrModel)},\n"
            + $"    vad_model={JsonSerializer.Serialize(diarization.VadModel)},\n"
            + $"    punc_model={JsonSerializer.Serialize(diarization.PunctuationModel)},\n"
            + $"    spk_model={JsonSerializer.Serialize(diarization.SpeakerModel)},\n"
            + "    device='cpu',\n"
            + "    disable_update=True,\n"
            + ")\n"
            + "res = model.generate(\n"
            + "    input='/workspace/input.wav',\n"
            + $"    batch_size_s={diarization.BatchSizeSeconds.ToString(System.Globalization.CultureInfo.InvariantCulture)},\n"
            + $"    batch_size_threshold_s={diarization.BatchSizeThresholdSeconds.ToString(System.Globalization.CultureInfo.InvariantCulture)}"
            + presetSpeakerCount
            + "\n)\n"
            + $"print('{JsonStart}')\n"
            + "print(json.dumps(res, ensure_ascii=False, default=str))\n"
            + $"print('{JsonEnd}')\n"
            + "PY\n"
            + "MODELSCOPE_CACHE=/workspace/models/modelscope python /tmp/meeting_assistant_diarize.py";
    }

    private static string ExtractJson(string output)
    {
        var start = output.IndexOf(JsonStart, StringComparison.Ordinal);
        if (start < 0)
        {
            throw new InvalidOperationException("FunASR diarization output did not contain the JSON start marker.");
        }

        start += JsonStart.Length;
        var end = output.IndexOf(JsonEnd, start, StringComparison.Ordinal);
        if (end < 0)
        {
            throw new InvalidOperationException("FunASR diarization output did not contain the JSON end marker.");
        }

        return output[start..end].Trim();
    }

    private static IReadOnlyList<TranscriptionSegment> ParseSegments(string json)
    {
        using var document = JsonDocument.Parse(json);
        var segments = new List<TranscriptionSegment>();

        foreach (var result in document.RootElement.EnumerateArray())
        {
            if (!result.TryGetProperty("sentence_info", out var sentences)
                || sentences.ValueKind != JsonValueKind.Array)
            {
                continue;
            }

            foreach (var sentence in sentences.EnumerateArray())
            {
                var text = ReadString(sentence, "text").Trim();
                if (string.IsNullOrWhiteSpace(text))
                {
                    continue;
                }

                segments.Add(new TranscriptionSegment(
                    TimeSpan.FromMilliseconds(ReadDouble(sentence, "start")),
                    TimeSpan.FromMilliseconds(ReadDouble(sentence, "end")),
                    $"Speaker {ReadInt(sentence, "spk")}",
                    text));
            }
        }

        return segments;
    }

    private static string ReadString(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property)
            && property.ValueKind == JsonValueKind.String
            ? property.GetString() ?? ""
            : "";
    }

    private static double ReadDouble(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property)
            && property.ValueKind == JsonValueKind.Number
            && property.TryGetDouble(out var value)
            ? value
            : 0;
    }

    private static int ReadInt(JsonElement element, string propertyName)
    {
        return element.TryGetProperty(propertyName, out var property)
            && property.ValueKind == JsonValueKind.Number
            && property.TryGetInt32(out var value)
            ? value
            : 0;
    }
}