using System.Net.Http.Json;
using System.Text.Json;

namespace Nop.Plugin.Misc.FruitBankPlugin.Services
{
    public class ReplicateService
    {
        private readonly HttpClient _http;

        public ReplicateService(HttpClient http)
        {
            _http = http;
        }
        public async Task<string> GenerateImageAsync(string prompt, bool removeBackground)
        {
            return await GenerateImageAsync("https://api.replicate.com/v1/models/black-forest-labs/flux-schnell/predictions", prompt, removeBackground);
        }
        public async Task<string> GenerateLogoAsync(string prompt, bool removeBackground)
        {
            return await GenerateImageAsync("https://api.replicate.com/v1/models/google/imagen-4-fast/predictions", prompt, removeBackground);            
        }

        public async Task<string> GenerateImageAsync(string apiUrl, string prompt, bool removeBackground)
        {
            var request = new
            {
                input = new { prompt = prompt, aspect_ratio = "1:1", output_format = "jpg" }
                
            };

            var createResponse = await _http.PostAsJsonAsync(apiUrl, request);
            createResponse.EnsureSuccessStatusCode();

            var createJson = await createResponse.Content.ReadFromJsonAsync<JsonElement>();
            if (!createJson.TryGetProperty("id", out var idProp))
                throw new Exception("Replicate response missing prediction ID.");

            string predictionId = idProp.GetString();
            string status = "";
            JsonElement finalJson;

            for (int attempt = 0; attempt < 30; attempt++)
            {
                var getResponse = await _http.GetAsync($"https://api.replicate.com/v1/predictions/{predictionId}");
                getResponse.EnsureSuccessStatusCode();

                finalJson = await getResponse.Content.ReadFromJsonAsync<JsonElement>();
                status = finalJson.GetProperty("status").GetString();

                if (status == "succeeded")
                {
                    var output = finalJson.GetProperty("output");

                    if (output.ValueKind == JsonValueKind.String)
                    {
                        var imageUrl = output.GetString();
                        if (removeBackground)
                        {
                            return await RemoveBackgroundAsync(imageUrl);

                        }
                        else
                        {
                            return imageUrl;
                        }
                    }
                    else if (output.ValueKind == JsonValueKind.Array)
                    {
                        var length = output.GetArrayLength();
                        var imageUrl = output[0].ToString();
                        if (removeBackground)
                        {
                            return await RemoveBackgroundAsync(imageUrl);

                        }
                        else
                        {
                            return imageUrl;
                        }
                    }


                    return "Replicate response succeeded but no output image URL found.";
                }
                else if (status == "failed")
                {
                    return "Replicate prediction failed.";
                }

                await Task.Delay(2500);
            }

            return "Timeout waiting for Replicate prediction to complete.";
        }

        public async Task<string> RemoveBackgroundAsync(string imageUrl)
        {
            var request = new
            {
                version = "a029dff38972b5fda4ec5d75d7d1cd25aeff621d2cf4946a41055d7db66b80bc",
                input = new { image = imageUrl }
            };

            var createResponse = await _http.PostAsJsonAsync("https://api.replicate.com/v1/predictions", request);
            createResponse.EnsureSuccessStatusCode();

            var createJson = await createResponse.Content.ReadFromJsonAsync<JsonElement>();
            if (!createJson.TryGetProperty("id", out var idProp))
                throw new Exception("Replicate response missing prediction ID.");

            string predictionId = idProp.GetString();
            string status = "";
            JsonElement finalJson;

            for (int attempt = 0; attempt < 20; attempt++)
            {
                var getResponse = await _http.GetAsync($"https://api.replicate.com/v1/predictions/{predictionId}");
                getResponse.EnsureSuccessStatusCode();

                finalJson = await getResponse.Content.ReadFromJsonAsync<JsonElement>();
                status = finalJson.GetProperty("status").GetString();

                if (status == "succeeded")
                {
                    var output = finalJson.GetProperty("output");

                    if (output.ValueKind == JsonValueKind.String)
                        return output.GetString();

                    return "Replicate response succeeded but no output image URL found.";
                }
                else if (status == "failed")
                {
                    return "Replicate prediction failed.";
                }

                await Task.Delay(1500);
            }

            return "Timeout waiting for Replicate prediction to complete.";
        }
        
        public async Task<JsonElement> AnalyzeImageAsync(string imageUrl)
        {
            var body = new
            {
                version = "a524caeaa23495bc9edc805ab08ab5fe943afd3febed884a4f3747aa32e9cd61",
                input = new
                {
                    image = imageUrl
                }
            };

            var request = new HttpRequestMessage(
                HttpMethod.Post,
                "https://api.replicate.com/v1/predictions")
            {
                Content = JsonContent.Create(body)
            };

            // This makes Replicate wait until the prediction is done
            request.Headers.Add("Prefer", "wait");

            var response = await _http.SendAsync(request);
            response.EnsureSuccessStatusCode();

            var json = await response.Content.ReadFromJsonAsync<JsonElement>();
            return json.ValueKind == JsonValueKind.Undefined
                ? throw new Exception("Empty response from Replicate.")
                : json;
        }
    }
}
