using System.Net;
using System.Net.Http.Json;
using System.Text.Json;
using Agent.Core.Tools;

namespace Agent.Core.Ollama;

public sealed class OllamaClient : IOllamaClient
{
    private static readonly TimeSpan[] RetryDelays = [TimeSpan.FromMilliseconds(250), TimeSpan.FromMilliseconds(500), TimeSpan.FromSeconds(1)];
    private readonly HttpClient _httpClient;
    private readonly OllamaSettings _settings;
    private readonly Func<TimeSpan, CancellationToken, Task> _delay;
    private bool? _visionCapability;

    public OllamaClient(HttpClient httpClient, OllamaSettings settings, Func<TimeSpan, CancellationToken, Task>? delay = null)
    {
        _httpClient = httpClient;
        _httpClient.BaseAddress = settings.Endpoint;
        _settings = settings;
        _delay = delay ?? Task.Delay;
    }

    public async Task<OllamaResult> ChatAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, CancellationToken cancellationToken)
        => await ChatCoreAsync(messages, temperatureOverride, null, cancellationToken).ConfigureAwait(false);

    public async Task<OllamaResult> ChatWithToolsAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, IReadOnlyList<InternalToolDefinition> tools, CancellationToken cancellationToken)
        => await ChatCoreAsync(messages, temperatureOverride, tools, cancellationToken).ConfigureAwait(false);

    public async Task<bool?> SupportsVisionAsync(CancellationToken cancellationToken)
    {
        if (_visionCapability is not null) return _visionCapability;
        try
        {
            using HttpResponseMessage response = await _httpClient.PostAsJsonAsync("api/show", new { model = _settings.Model }, cancellationToken).ConfigureAwait(false);
            if (!response.IsSuccessStatusCode) return null;
            using JsonDocument document = await JsonDocument.ParseAsync(await response.Content.ReadAsStreamAsync(cancellationToken).ConfigureAwait(false), cancellationToken: cancellationToken).ConfigureAwait(false);
            if (!document.RootElement.TryGetProperty("capabilities", out JsonElement capabilities) || capabilities.ValueKind != JsonValueKind.Array) return null;
            _visionCapability = capabilities.EnumerateArray().Any(static item => item.GetString()?.Equals("vision", StringComparison.OrdinalIgnoreCase) == true);
            return _visionCapability;
        }
        catch (HttpRequestException) { return null; }
    }

    private async Task<OllamaResult> ChatCoreAsync(IReadOnlyList<OllamaMessage> messages, double? temperatureOverride, IReadOnlyList<InternalToolDefinition>? tools, CancellationToken cancellationToken)
    {
        Dictionary<string, JsonElement> options = _settings.Options.ToDictionary(static pair => pair.Key, static pair => pair.Value, StringComparer.Ordinal);
        if (temperatureOverride is not null) options["temperature"] = JsonSerializer.SerializeToElement(temperatureOverride.Value);
        var request = new OllamaChatRequest
        {
            Model = _settings.Model,
            Messages = messages,
            Stream = false,
            Think = _settings.Think,
            KeepAlive = _settings.KeepAlive,
            Options = options,
            Tools = tools?.Select(static tool => new OllamaToolDefinition("function", new OllamaFunctionDefinition(tool.Name, tool.Description, tool.Parameters))).ToArray()
        };

        for (int attempt = 0; ; attempt++)
        {
            try
            {
                using HttpResponseMessage response = await _httpClient.PostAsJsonAsync("api/chat", request, cancellationToken).ConfigureAwait(false);
                if (!response.IsSuccessStatusCode)
                {
                    string body = await response.Content.ReadAsStringAsync(cancellationToken).ConfigureAwait(false);
                    bool transient = IsTransient(response.StatusCode);
                    if (transient && attempt < RetryDelays.Length)
                    {
                        await _delay(RetryDelays[attempt], cancellationToken).ConfigureAwait(false);
                        continue;
                    }
                    throw new OllamaException($"Ollama returned HTTP {(int)response.StatusCode}: {body}", transient);
                }

                OllamaChatResponse result = await response.Content.ReadFromJsonAsync<OllamaChatResponse>(cancellationToken: cancellationToken).ConfigureAwait(false)
                    ?? throw new OllamaException("Ollama returned an empty response.", false);
                if (!string.IsNullOrWhiteSpace(result.Error)) throw new OllamaException(result.Error, false);
                if (result.Message is null) throw new OllamaException("Ollama response has no message.", false);
                return new OllamaResult(result.Message.Content, result.PromptEvalCount, result.EvalCount, result.Message.ToolCalls ?? []);
            }
            catch (Exception ex) when (ex is HttpRequestException or TaskCanceledException && !cancellationToken.IsCancellationRequested)
            {
                if (attempt >= RetryDelays.Length) throw new OllamaException("Ollama is unavailable after three retries.", true, ex);
                await _delay(RetryDelays[attempt], cancellationToken).ConfigureAwait(false);
            }
        }
    }

    private static bool IsTransient(HttpStatusCode statusCode) => statusCode is HttpStatusCode.RequestTimeout or HttpStatusCode.TooManyRequests || (int)statusCode >= 500;
}