using Newtonsoft.Json.Linq;
using System;
using System.Collections.Generic;
using System.Net.Http;
using System.Threading;
using System.Threading.Tasks;
using UnityEngine;

public class TrellisClient : MonoBehaviour
{
    public static TrellisClient Instance { get; private set; }

    public string TRELLIS_BASE_URL;

    private HttpClient httpClient;

    private void Awake()
    {
        httpClient = new HttpClient();
        httpClient.BaseAddress = new Uri(TRELLIS_BASE_URL);

        Instance = this;
    }

    // Start is called before the first frame update
    void Start()
    {
        
    }

    // Update is called once per frame
    void Update()
    {
        
    }


    public async Task<byte[]> GenerateModel(
            string imageBase64,
            int seed = 42,
            int pollIntervalMs = 1000)
    {
        // --- Set generation parameters (form-encoded, like Python requests.post(data=params)) ---
        var form = new Dictionary<string, string>
        {
            ["image_base64"] = imageBase64,
            ["seed"] = seed.ToString(),
            ["ss_guidance_strength"] = "7.5",
            ["ss_sampling_steps"] = "10",
            ["slat_guidance_strength"] = "7.5",
            ["slat_sampling_steps"] = "10",
            ["mesh_simplify_ratio"] = "0.99",
            ["texture_size"] = "1024",
            ["texture_opt_total_steps"] = "1000",
            ["output_format"] = "glb"
        };

        Debug.Log("Starting Trellis generation...");
        using (var content = new FormUrlEncodedContent(form))
        using (var startResp = await httpClient.PostAsync("generate_no_preview", content))
        {
            startResp.EnsureSuccessStatusCode();
        }

        // --- Poll /status until COMPLETE or FAILED ---
        while (true)
        {
            using (var statusResp = await httpClient.GetAsync("status"))
            {
                statusResp.EnsureSuccessStatusCode();
                var json = await statusResp.Content.ReadAsStringAsync();
                var status = JObject.Parse(json);

                var progress = status.Value<int?>("progress") ?? 0;
                var state = status.Value<string>("status") ?? "UNKNOWN";
                Debug.Log($"TRELLIS progress: {progress}% (status: {state})");

                if (string.Equals(state, "COMPLETE", StringComparison.OrdinalIgnoreCase))
                    break;

                if (string.Equals(state, "FAILED", StringComparison.OrdinalIgnoreCase))
                {
                    var msg = status.Value<string>("message") ?? "Generation failed.";
                    throw new InvalidOperationException($"Trellis generation failed: {msg}");
                }
            }

            await Task.Delay(pollIntervalMs);
        }

        // --- Download the model binary ---
        Debug.Log("Downloading model...");
        using (var downloadResponse = await httpClient.GetAsync("download/model"))
        {
            downloadResponse.EnsureSuccessStatusCode();
            var bytes = await downloadResponse.Content.ReadAsByteArrayAsync();
            Debug.Log($"Downloaded {bytes.Length} bytes");
            return bytes;
        }
    }

}