using System.Text.Json;

namespace BLAIzor.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.";
        }
    }
}