using System;
using System.Collections.Generic;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;
using UnityEngine.Networking;

namespace CG.Chat
{
    public sealed class OllamaChatProvider : IChatProvider
    {
        private readonly string _baseUrl;

        public OllamaChatProvider(string baseUrl)
        {
            _baseUrl = (baseUrl ?? "").TrimEnd('/');
        }

        [Serializable] private class ReqMsg { public string role; public string content; }
        [Serializable] private class ChatReq { public string model; public ReqMsg[] messages; public bool stream = false; }
        [Serializable] private class RespMsg { public string role; public string content; }
        [Serializable] private class ChatResp { public RespMsg message; }

        public async Task<string> SendChatAsync(IReadOnlyList<ChatMessage> messages, string model, CancellationToken ct)
        {
            var url = $"{_baseUrl}/api/chat";

            var req = new ChatReq
            {
                model = model,
                stream = false,
                messages = Build(messages)
            };

            var json = JsonUtility.ToJson(req);

            using var uwr = new UnityWebRequest(url, "POST");
            uwr.uploadHandler = new UploadHandlerRaw(Encoding.UTF8.GetBytes(json));
            uwr.downloadHandler = new DownloadHandlerBuffer();
            uwr.SetRequestHeader("Content-Type", "application/json");

            var op = uwr.SendWebRequest();
            while (!op.isDone)
            {
                if (ct.IsCancellationRequested)
                {
                    uwr.Abort();
                    ct.ThrowIfCancellationRequested();
                }
                await Task.Yield();
            }

            if (uwr.result != UnityWebRequest.Result.Success)
                throw new Exception($"Ollama request failed: {uwr.error}\n{uwr.downloadHandler.text}");

            var text = uwr.downloadHandler.text;
            var resp = JsonUtility.FromJson<ChatResp>(text);
            return resp?.message?.content ?? "";
        }

        private static ReqMsg[] Build(IReadOnlyList<ChatMessage> messages)
        {
            var arr = new ReqMsg[messages.Count];
            for (int i = 0; i < messages.Count; i++)
                arr[i] = new ReqMsg { role = messages[i].role, content = messages[i].content };
            return arr;
        }
    }
}
