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

namespace MeetingAssistant.Transcription;

public sealed class PyannoteTranscriptFinalizer
{
    private const string JsonStart = "__MEETING_ASSISTANT_PYANNOTE_JSON_START__";
    private const string JsonEnd = "__MEETING_ASSISTANT_PYANNOTE_JSON_END__";

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

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

    public async Task<IReadOnlyList<TranscriptionSegment>> FinalizeAsync(
        string audioPath,
        IReadOnlyList<TranscriptionSegment> liveSegments,
        SpeechRecognitionPipelineOptions pipelineOptions,
        CancellationToken cancellationToken)
    {
        return await FinalizeAsync(
            audioPath,
            liveSegments,
            options.WhisperLocal.Diarization,
            pipelineOptions,
            cancellationToken);
    }

    public async Task<IReadOnlyList<TranscriptionSegment>> FinalizeAsync(
        string audioPath,
        IReadOnlyList<TranscriptionSegment> liveSegments,
        PyannoteDiarizationOptions diarization,
        SpeechRecognitionPipelineOptions pipelineOptions,
        CancellationToken cancellationToken)
    {
        if (!diarization.Enabled || liveSegments.Count == 0)
        {
            return [];
        }

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

        var token = ResolveToken(diarization);
        if (string.IsNullOrWhiteSpace(token))
        {
            logger.LogWarning(
                "Pyannote diarization is enabled but no Hugging Face token is configured directly or via {TokenEnv}",
                diarization.TokenEnv);
            return [];
        }

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

        var turns = ParseTurns(ExtractJson(result.StandardOutput));
        var segments = diarization.AlignmentMode == PyannoteAlignmentMode.PyannoteTurns
            ? BuildTurnSegments(liveSegments, turns)
            : AssignSpeakers(liveSegments, turns);
        logger.LogInformation(
            "pyannote diarization produced {TurnCount} speaker turns and built {SegmentCount} transcript segments using {AlignmentMode}",
            turns.Count,
            segments.Count,
            diarization.AlignmentMode);
        return segments;
    }

    public async Task WarmUpAsync(
        PyannoteDiarizationOptions diarization,
        CancellationToken cancellationToken)
    {
        if (!diarization.Enabled)
        {
            return;
        }

        var token = ResolveToken(diarization);
        if (string.IsNullOrWhiteSpace(token))
        {
            logger.LogWarning(
                "Pyannote warm-up is enabled but no Hugging Face token is configured directly or via {TokenEnv}",
                diarization.TokenEnv);
            return;
        }

        var modelsFolder = VaultPath.Resolve(diarization.ModelsFolder);
        Directory.CreateDirectory(modelsFolder);
        using var timeoutSource = diarization.CommandTimeout > TimeSpan.Zero
            ? new CancellationTokenSource(diarization.CommandTimeout)
            : null;
        using var linkedSource = timeoutSource is null
            ? null
            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);

        try
        {
            if (diarization.BuildImage)
            {
                await EnsureDockerImageAsync(diarization, modelsFolder, linkedSource?.Token ?? cancellationToken);
            }

            var result = await commandRunner.RunAsync(
                diarization.DockerCommand,
                BuildWarmUpDockerArguments(diarization, modelsFolder),
                linkedSource?.Token ?? cancellationToken,
                new Dictionary<string, string> { ["HF_TOKEN"] = token });
            if (result.ExitCode != 0)
            {
                throw new InvalidOperationException(
                    $"pyannote warm-up failed with exit code {result.ExitCode}: {result.StandardError}");
            }
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
        {
            throw new TimeoutException(
                $"pyannote warm-up timed out after {diarization.CommandTimeout}.");
        }
    }

    private async Task<CommandResult> RunDiarizationAsync(
        string fullAudioPath,
        string token,
        PyannoteDiarizationOptions diarization,
        SpeechRecognitionPipelineOptions pipelineOptions,
        CancellationToken cancellationToken)
    {
        var modelsFolder = VaultPath.Resolve(diarization.ModelsFolder);
        Directory.CreateDirectory(modelsFolder);
        using var timeoutSource = diarization.CommandTimeout > TimeSpan.Zero
            ? new CancellationTokenSource(diarization.CommandTimeout)
            : null;
        using var linkedSource = timeoutSource is null
            ? null
            : CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutSource.Token);

        try
        {
            if (diarization.BuildImage)
            {
                await EnsureDockerImageAsync(diarization, modelsFolder, linkedSource?.Token ?? cancellationToken);
            }

            return await commandRunner.RunAsync(
                diarization.DockerCommand,
                BuildDockerArguments(diarization, fullAudioPath, modelsFolder, pipelineOptions),
                linkedSource?.Token ?? cancellationToken,
                new Dictionary<string, string> { ["HF_TOKEN"] = token });
        }
        catch (OperationCanceledException) when (!cancellationToken.IsCancellationRequested && timeoutSource?.IsCancellationRequested == true)
        {
            throw new TimeoutException(
                $"pyannote diarization timed out after {diarization.CommandTimeout}.");
        }
    }

    private async Task EnsureDockerImageAsync(
        PyannoteDiarizationOptions diarization,
        string modelsFolder,
        CancellationToken cancellationToken)
    {
        var inspectResult = await commandRunner.RunAsync(
            diarization.DockerCommand,
            ["image", "inspect", "--format", "{{.Id}}", diarization.Image],
            cancellationToken);
        if (inspectResult.ExitCode == 0)
        {
            return;
        }

        var dockerfilePath = Path.Combine(modelsFolder, "docker", $"{SanitizeFileName(diarization.Image)}.Dockerfile");
        Directory.CreateDirectory(Path.GetDirectoryName(dockerfilePath)!);
        await File.WriteAllTextAsync(dockerfilePath, BuildDockerfile(diarization), cancellationToken);
        var result = await commandRunner.RunAsync(
            diarization.DockerCommand,
            ["build", "-t", diarization.Image, "-f", dockerfilePath, Path.GetDirectoryName(dockerfilePath)!],
            cancellationToken);
        if (result.ExitCode != 0)
        {
            throw new InvalidOperationException(
                $"pyannote Docker image build failed with exit code {result.ExitCode}: {result.StandardError}");
        }
    }

    private static string BuildDockerfile(PyannoteDiarizationOptions diarization)
    {
        return
            $"FROM {diarization.BaseImage}\n"
            + "ENV DEBIAN_FRONTEND=noninteractive\n"
            + "RUN apt-get update \\\n"
            + "    && apt-get install -y --no-install-recommends ffmpeg libsndfile1 \\\n"
            + "    && rm -rf /var/lib/apt/lists/*\n"
            + "RUN python -m pip install --upgrade pip \\\n"
            + "    && python -m pip install pyannote.audio soundfile\n";
    }

    private string[] BuildDockerArguments(
        PyannoteDiarizationOptions diarization,
        string fullAudioPath,
        string modelsFolder,
        SpeechRecognitionPipelineOptions pipelineOptions)
    {
        return
        [
            "run",
            "--rm",
            "-e",
            "HF_TOKEN",
            "-v",
            $"{fullAudioPath}:/workspace/input.wav:ro",
            "-v",
            $"{modelsFolder}:/workspace/cache",
            "-e",
            "HF_HOME=/workspace/cache/huggingface",
            "-e",
            "XDG_CACHE_HOME=/workspace/cache",
            "-e",
            "PIP_CACHE_DIR=/workspace/cache/pip",
            "-e",
            "TORCH_HOME=/workspace/cache/torch",
            diarization.Image,
            "sh",
            "-lc",
            BuildPythonCommand(diarization, pipelineOptions)
        ];
    }

    private string[] BuildWarmUpDockerArguments(
        PyannoteDiarizationOptions diarization,
        string modelsFolder)
    {
        return
        [
            "run",
            "--rm",
            "-e",
            "HF_TOKEN",
            "-v",
            $"{modelsFolder}:/workspace/cache",
            "-e",
            "HF_HOME=/workspace/cache/huggingface",
            "-e",
            "XDG_CACHE_HOME=/workspace/cache",
            "-e",
            "PIP_CACHE_DIR=/workspace/cache/pip",
            "-e",
            "TORCH_HOME=/workspace/cache/torch",
            diarization.Image,
            "sh",
            "-lc",
            BuildWarmUpPythonCommand(diarization)
        ];
    }

    private static string BuildPythonCommand(
        PyannoteDiarizationOptions diarization,
        SpeechRecognitionPipelineOptions pipelineOptions)
    {
        var model = diarization.Model;
        var annotationProperty = diarization.AnnotationSource == PyannoteAnnotationSource.ExclusiveSpeakerDiarization
            ? "exclusive_speaker_diarization"
            : "speaker_diarization";
        return
            "cat > /tmp/meeting_assistant_pyannote.py <<'PY'\n"
            + "import json\n"
            + "import os\n"
            + "import torch\n"
            + "from pyannote.audio import Pipeline\n"
            + $"pipeline = Pipeline.from_pretrained({JsonSerializer.Serialize(model)}, token=os.environ.get('HF_TOKEN'))\n"
            + "pipeline.to(torch.device('cpu'))\n"
            + BuildPipelineInvocation(pipelineOptions)
            + $"diarization = getattr(result, {JsonSerializer.Serialize(annotationProperty)}, result)\n"
            + "turns = []\n"
            + "for turn, _, speaker in diarization.itertracks(yield_label=True):\n"
            + "    turns.append({'start': float(turn.start), 'end': float(turn.end), 'speaker': str(speaker)})\n"
            + $"print('{JsonStart}')\n"
            + "print(json.dumps(turns, ensure_ascii=False))\n"
            + $"print('{JsonEnd}')\n"
            + "PY\n"
            + "python /tmp/meeting_assistant_pyannote.py";
    }

    private static string BuildWarmUpPythonCommand(PyannoteDiarizationOptions diarization)
    {
        var model = diarization.Model;
        return
            "cat > /tmp/meeting_assistant_pyannote_warmup.py <<'PY'\n"
            + "import os\n"
            + "import torch\n"
            + "from pyannote.audio import Pipeline\n"
            + $"pipeline = Pipeline.from_pretrained({JsonSerializer.Serialize(model)}, token=os.environ.get('HF_TOKEN'))\n"
            + "pipeline.to(torch.device('cpu'))\n"
            + "print('pyannote warm-up complete')\n"
            + "PY\n"
            + "python /tmp/meeting_assistant_pyannote_warmup.py";
    }

    private static string BuildPipelineInvocation(SpeechRecognitionPipelineOptions pipelineOptions)
    {
        return pipelineOptions.NumSpeakers is > 0
            ? $"result = pipeline('/workspace/input.wav', num_speakers={pipelineOptions.NumSpeakers.Value})\n"
            : "result = pipeline('/workspace/input.wav')\n";
    }

    private static IReadOnlyList<TranscriptionSegment> BuildTurnSegments(
        IReadOnlyList<TranscriptionSegment> liveSegments,
        IReadOnlyList<SpeakerTurn> turns)
    {
        if (turns.Count == 0)
        {
            return [];
        }

        var segments = new List<TranscriptionSegment>();
        foreach (var turn in turns)
        {
            var text = BuildTurnText(liveSegments, turn);
            if (string.IsNullOrWhiteSpace(text))
            {
                continue;
            }

            segments.Add(new TranscriptionSegment(
                TimeSpan.FromSeconds(turn.Start),
                TimeSpan.FromSeconds(turn.End),
                turn.Speaker,
                text));
        }

        return segments.Count == 0
            ? AssignSpeakers(liveSegments, turns)
            : segments;
    }

    private static string BuildTurnText(IReadOnlyList<TranscriptionSegment> liveSegments, SpeakerTurn turn)
    {
        var parts = new List<string>();
        foreach (var segment in liveSegments)
        {
            var segmentStart = segment.Start.TotalSeconds;
            var segmentEnd = segment.End > segment.Start
                ? segment.End.TotalSeconds
                : segmentStart;
            var overlapStart = Math.Max(segmentStart, turn.Start);
            var overlapEnd = Math.Min(segmentEnd, turn.End);
            if (overlapEnd <= overlapStart)
            {
                continue;
            }

            var part = SliceTextByTime(segment.Text, segmentStart, segmentEnd, overlapStart, overlapEnd);
            if (!string.IsNullOrWhiteSpace(part))
            {
                parts.Add(part);
            }
        }

        return string.Join(' ', parts);
    }

    private static string SliceTextByTime(
        string text,
        double segmentStart,
        double segmentEnd,
        double overlapStart,
        double overlapEnd)
    {
        var words = text
            .Split(' ', StringSplitOptions.RemoveEmptyEntries | StringSplitOptions.TrimEntries);
        if (words.Length == 0)
        {
            return "";
        }

        var duration = Math.Max(0.001, segmentEnd - segmentStart);
        var startFraction = Math.Clamp((overlapStart - segmentStart) / duration, 0, 1);
        var endFraction = Math.Clamp((overlapEnd - segmentStart) / duration, 0, 1);
        var startIndex = Math.Clamp((int)Math.Floor(startFraction * words.Length), 0, words.Length - 1);
        var endIndex = Math.Clamp((int)Math.Ceiling(endFraction * words.Length), startIndex + 1, words.Length);
        return string.Join(' ', words[startIndex..endIndex]);
    }

    private static IReadOnlyList<TranscriptionSegment> AssignSpeakers(
        IReadOnlyList<TranscriptionSegment> liveSegments,
        IReadOnlyList<SpeakerTurn> turns)
    {
        if (turns.Count == 0)
        {
            return [];
        }

        return liveSegments
            .Select(segment =>
            {
                var speaker = FindBestSpeaker(segment, turns) ?? segment.Speaker;
                return segment with { Speaker = speaker };
            })
            .ToList();
    }

    private static string? FindBestSpeaker(TranscriptionSegment segment, IReadOnlyList<SpeakerTurn> turns)
    {
        var segmentStart = segment.Start.TotalSeconds;
        var segmentEnd = segment.End > segment.Start
            ? segment.End.TotalSeconds
            : segmentStart;
        var bestOverlap = 0d;
        string? bestSpeaker = null;

        foreach (var turn in turns)
        {
            var overlap = Math.Min(segmentEnd, turn.End) - Math.Max(segmentStart, turn.Start);
            if (overlap > bestOverlap)
            {
                bestOverlap = overlap;
                bestSpeaker = turn.Speaker;
            }
        }

        return bestSpeaker;
    }

    private static string ExtractJson(string output)
    {
        var start = output.IndexOf(JsonStart, StringComparison.Ordinal);
        if (start < 0)
        {
            throw new InvalidOperationException("pyannote 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("pyannote output did not contain the JSON end marker.");
        }

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

    private static IReadOnlyList<SpeakerTurn> ParseTurns(string json)
    {
        using var document = JsonDocument.Parse(json);
        var turns = new List<SpeakerTurn>();
        foreach (var item in document.RootElement.EnumerateArray())
        {
            var speaker = ReadString(item, "speaker");
            if (string.IsNullOrWhiteSpace(speaker))
            {
                continue;
            }

            turns.Add(new SpeakerTurn(
                ReadDouble(item, "start"),
                ReadDouble(item, "end"),
                speaker));
        }

        return turns;
    }

    private static string ResolveToken(PyannoteDiarizationOptions options)
    {
        if (!string.IsNullOrWhiteSpace(options.Token))
        {
            return options.Token;
        }

        return string.IsNullOrWhiteSpace(options.TokenEnv)
            ? ""
            : Environment.GetEnvironmentVariable(options.TokenEnv) ?? "";
    }

    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 string SanitizeFileName(string value)
    {
        var invalid = Path.GetInvalidFileNameChars();
        return string.Concat(value.Select(character => invalid.Contains(character) ? '_' : character));
    }

    private sealed record SpeakerTurn(double Start, double End, string Speaker);
}