using System.Net.Http.Headers;
using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using OpenQuery.Models;

namespace OpenQuery.Services;

public class OpenRouterClient
{
    private readonly HttpClient _httpClient;
    private readonly string _apiKey;
    private readonly string _baseUrl = "https://openrouter.ai/api/v1";

    public OpenRouterClient(string apiKey)
    {
        _apiKey = apiKey;
        _httpClient = new HttpClient();
        _httpClient.DefaultRequestHeaders.Authorization = new AuthenticationHeaderValue("Bearer", apiKey);
        _httpClient.DefaultRequestHeaders.Accept.Add(new MediaTypeWithQualityHeaderValue("application/json"));
    }

    public async IAsyncEnumerable<StreamChunk> StreamAsync(ChatCompletionRequest request, [EnumeratorCancellation] CancellationToken cancellationToken = default)
    {
        request = request with { Stream = true };

        var json = JsonSerializer.Serialize(request, AppJsonContext.Default.ChatCompletionRequest);
        var content = new StringContent(json, Encoding.UTF8, new MediaTypeHeaderValue("application/json"));

        var httpRequest = new HttpRequestMessage(HttpMethod.Post, $"{_baseUrl}/chat/completions")
        {
            Content = content
        };

        using var response = await _httpClient.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead, cancellationToken);
        response.EnsureSuccessStatusCode();

        await using var stream = await response.Content.ReadAsStreamAsync(cancellationToken);
        using var reader = new StreamReader(stream);

        while (await reader.ReadLineAsync(cancellationToken) is { } line)
        {
            if (string.IsNullOrEmpty(line) || line.StartsWith($":"))
                continue;

            if (!line.StartsWith("data: ")) continue;
            var data = line[6..];
            if (data == "[DONE]")
                yield break;

            var chunk = JsonSerializer.Deserialize<ChatCompletionChunk>(data, AppJsonContext.Default.ChatCompletionChunk);
            if (!(chunk?.Choices?.Count > 0)) continue;
            var delta = chunk.Choices[0].Delta;
            if (!string.IsNullOrEmpty(delta?.Content))
                yield return new StreamChunk(delta.Content);

            if (delta?.ToolCalls is not { Count: > 0 }) continue;
            var toolCall = delta.ToolCalls[0];
            yield return new StreamChunk(null, new ClientToolCall(
                toolCall.Id,
                toolCall.Function.Name,
                toolCall.Function.Arguments
            ));
        }
    }

    public async Task<ChatCompletionResponse> CompleteAsync(ChatCompletionRequest request)
    {
        request = request with { Stream = false };

        var json = JsonSerializer.Serialize(request, AppJsonContext.Default.ChatCompletionRequest);
        var content = new StringContent(json, Encoding.UTF8, new MediaTypeHeaderValue("application/json"));

        var response = await _httpClient.PostAsync($"{_baseUrl}/chat/completions", content);
        response.EnsureSuccessStatusCode();

        var responseJson = await response.Content.ReadAsStringAsync();
        return JsonSerializer.Deserialize<ChatCompletionResponse>(responseJson, AppJsonContext.Default.ChatCompletionResponse)!;
    }

    public async Task<float[][]> EmbedAsync(string model, List<string> inputs)
    {
        var request = new EmbeddingRequest(model, inputs);
        var json = JsonSerializer.Serialize(request, AppJsonContext.Default.EmbeddingRequest);
        var content = new StringContent(json, Encoding.UTF8, new MediaTypeHeaderValue("application/json"));

        var response = await _httpClient.PostAsync($"{_baseUrl}/embeddings", content);
        response.EnsureSuccessStatusCode();

        var responseJson = await response.Content.ReadAsStringAsync();
        var embeddingResponse = JsonSerializer.Deserialize<EmbeddingResponse>(responseJson, AppJsonContext.Default.EmbeddingResponse)!;

        return embeddingResponse.Data
            .OrderBy(d => d.Index)
            .Select(d => d.Embedding)
            .ToArray();
    }
}

public record StreamChunk(
    string? TextDelta = null,
    ClientToolCall? Tool = null
);

public record ClientToolCall(
    string ToolId,
    string ToolName,
    string Arguments
);

public record ChatCompletionChunk(
    [property: JsonPropertyName("choices")] List<ChunkChoice> Choices
);

public record ChunkChoice(
    [property: JsonPropertyName("delta")] ChunkDelta Delta
);

public record ChunkDelta(
    [property: JsonPropertyName("content")] string? Content = null,
    [property: JsonPropertyName("tool_calls")] List<ToolCall>? ToolCalls = null
);