// PUBLIC DOMAIN - NO LICENSE, NO WARRANTY
// Copyright 2025 TimeHexOn & foxhop & russell@unturf
// https://www.permacomputer.com

const std = @import("std");
const http = std.http;
const json = std.json;

// Model registry entry
pub const ModelInfo = struct {
    id: []const u8,
    endpoint: []const u8,
    max_tokens: u32,
};

// Chat completion structures
pub const ChatMessage = struct {
    role: []const u8,
    content: []const u8,
};

pub const ChatRequest = struct {
    model: []const u8,
    messages: []const ChatMessage,
    max_tokens: u32 = 1000,
    stream: bool = false,
};

// TTS request structure
pub const TTSRequest = struct {
    model: []const u8,
    voice: []const u8,
    input: []const u8,
};

// Streaming chunk response
pub const StreamChunk = struct {
    content: ?[]const u8,
    done: bool,

    pub fn deinit(self: *StreamChunk, allocator: std.mem.Allocator) void {
        if (self.content) |c| {
            allocator.free(c);
        }
    }
};

// Main client for UncloseAI API interactions
pub const Client = struct {
    allocator: std.mem.Allocator,
    models: std.ArrayList(ModelInfo),
    tts_endpoints: std.ArrayList([]const u8),

    pub fn init(allocator: std.mem.Allocator) !Client {
        var client = Client{
            .allocator = allocator,
            .models = std.ArrayList(ModelInfo).init(allocator),
            .tts_endpoints = std.ArrayList([]const u8).init(allocator),
        };

        try client.discoverModels();
        return client;
    }

    pub fn deinit(self: *Client) void {
        for (self.models.items) |model| {
            self.allocator.free(model.id);
            self.allocator.free(model.endpoint);
        }
        self.models.deinit();

        for (self.tts_endpoints.items) |endpoint| {
            self.allocator.free(endpoint);
        }
        self.tts_endpoints.deinit();
    }

    pub fn getModels(self: *Client) []const ModelInfo {
        return self.models.items;
    }

    pub fn getTTSEndpoints(self: *Client) []const []const u8 {
        return self.tts_endpoints.items;
    }

    fn discoverModels(self: *Client) !void {
        std.debug.print("Discovering models from environment variables...\n", .{});

        // Discover chat/code models from MODEL_ENDPOINT_1..9999
        var i: u32 = 1;
        while (i < 10000) : (i += 1) {
            var env_key_buf: [32]u8 = undefined;
            const env_key = try std.fmt.bufPrint(&env_key_buf, "MODEL_ENDPOINT_{d}", .{i});

            const endpoint = getEnv(env_key) orelse break;

            std.debug.print("Discovering from: {s}\n", .{endpoint});

            // Build /models URL
            var url_buf: [256]u8 = undefined;
            const models_url = try std.fmt.bufPrint(&url_buf, "{s}/models", .{endpoint});

            // Fetch models list
            const response = makeGETRequest(self.allocator, models_url) catch |err| {
                std.debug.print("Failed to fetch models from {s}: {}\n", .{ endpoint, err });
                continue;
            };
            defer self.allocator.free(response);

            // Parse JSON response - look for model IDs
            var pos: usize = 0;
            while (std.mem.indexOf(u8, response[pos..], "\"id\":\"")) |idx| {
                pos += idx + 6;
                const end = std.mem.indexOf(u8, response[pos..], "\"") orelse break;
                const model_id = try self.allocator.dupe(u8, response[pos .. pos + end]);

                const max_tokens: u32 = 8192;
                const endpoint_copy = try self.allocator.dupe(u8, endpoint);

                try self.models.append(ModelInfo{
                    .id = model_id,
                    .endpoint = endpoint_copy,
                    .max_tokens = max_tokens,
                });

                pos += end + 1;
            }
        }

        // Discover TTS endpoints from TTS_ENDPOINT_1..9999
        i = 1;
        while (i < 10000) : (i += 1) {
            var env_key_buf: [32]u8 = undefined;
            const env_key = try std.fmt.bufPrint(&env_key_buf, "TTS_ENDPOINT_{d}", .{i});

            const endpoint = getEnv(env_key) orelse break;

            std.debug.print("Discovering TTS from: {s}\n", .{endpoint});

            const endpoint_copy = try self.allocator.dupe(u8, endpoint);
            try self.tts_endpoints.append(endpoint_copy);
        }

        std.debug.print("\nDiscovered {d} model(s) and {d} TTS endpoint(s)\n", .{ self.models.items.len, self.tts_endpoints.items.len });
    }

    // Non-streaming chat completion
    pub fn chatCompletion(self: *Client, request: ChatRequest) ![]u8 {
        // Find the model's endpoint
        var model_endpoint: ?[]const u8 = null;
        for (self.models.items) |model| {
            if (std.mem.eql(u8, model.id, request.model)) {
                model_endpoint = model.endpoint;
                break;
            }
        }

        if (model_endpoint == null) {
            return error.ModelNotFound;
        }

        const json_string = try json.stringifyAlloc(self.allocator, request, .{});
        defer self.allocator.free(json_string);

        var url_buf: [512]u8 = undefined;
        const url = try std.fmt.bufPrint(&url_buf, "{s}/chat/completions", .{model_endpoint.?});

        return try makeHTTPRequest(self.allocator, url, json_string);
    }

    // Streaming chat completion - returns iterator-like stream
    pub fn chatCompletionStream(self: *Client, request: ChatRequest) !ChatStream {
        var model_endpoint: ?[]const u8 = null;
        for (self.models.items) |model| {
            if (std.mem.eql(u8, model.id, request.model)) {
                model_endpoint = model.endpoint;
                break;
            }
        }

        if (model_endpoint == null) {
            return error.ModelNotFound;
        }

        // Enable streaming in request
        var stream_request = request;
        stream_request.stream = true;

        return ChatStream.init(self.allocator, model_endpoint.?, stream_request);
    }

    // TTS generation
    pub fn generateSpeech(self: *Client, request: TTSRequest, tts_endpoint_idx: usize) ![]u8 {
        if (tts_endpoint_idx >= self.tts_endpoints.items.len) {
            return error.TTSEndpointNotFound;
        }

        const tts_endpoint = self.tts_endpoints.items[tts_endpoint_idx];

        const json_string = try json.stringifyAlloc(self.allocator, request, .{});
        defer self.allocator.free(json_string);

        var url_buf: [512]u8 = undefined;
        const url = try std.fmt.bufPrint(&url_buf, "{s}/audio/speech", .{tts_endpoint});

        return try makeHTTPRequest(self.allocator, url, json_string);
    }
};

// Chat streaming iterator
pub const ChatStream = struct {
    allocator: std.mem.Allocator,
    endpoint: []const u8,
    request: ChatRequest,
    client: ?http.Client = null,
    req: ?http.Client.Request = null,
    buffer: [4096]u8 = undefined,
    buffer_pos: usize = 0,
    buffer_len: usize = 0,
    done: bool = false,

    pub fn init(allocator: std.mem.Allocator, endpoint: []const u8, request: ChatRequest) !ChatStream {
        return ChatStream{
            .allocator = allocator,
            .endpoint = endpoint,
            .request = request,
        };
    }

    pub fn start(self: *ChatStream) !void {
        self.client = http.Client{ .allocator = self.allocator };

        const json_string = try json.stringifyAlloc(self.allocator, self.request, .{});
        defer self.allocator.free(json_string);

        var url_buf: [512]u8 = undefined;
        const url = try std.fmt.bufPrint(&url_buf, "{s}/chat/completions", .{self.endpoint});
        const full_uri = try std.Uri.parse(url);

        const server_header_buffer = try self.allocator.alloc(u8, 1024 * 8);
        errdefer self.allocator.free(server_header_buffer);

        self.req = try self.client.?.open(.POST, full_uri, .{
            .server_header_buffer = server_header_buffer,
            .extra_headers = &[_]http.Header{
                .{ .name = "Content-Type", .value = "application/json" },
                .{ .name = "Accept", .value = "text/event-stream" },
            },
        });

        self.req.?.transfer_encoding = .chunked;
        try self.req.?.send();
        try self.req.?.writeAll(json_string);
        try self.req.?.finish();
        try self.req.?.wait();
    }

    pub fn next(self: *ChatStream) !?StreamChunk {
        if (self.done) return null;

        // Read from HTTP response stream
        if (self.buffer_pos >= self.buffer_len) {
            // Need more data
            const n = try self.req.?.reader().read(&self.buffer);
            if (n == 0) {
                self.done = true;
                return StreamChunk{ .content = null, .done = true };
            }
            self.buffer_len = n;
            self.buffer_pos = 0;
        }

        // Parse SSE format: data: {...}\n\n
        const line_start = self.buffer_pos;
        var line_end = line_start;

        // Find end of line
        while (line_end < self.buffer_len and self.buffer[line_end] != '\n') {
            line_end += 1;
        }

        if (line_end >= self.buffer_len) {
            // Need more data for complete line
            return null;
        }

        const line = self.buffer[line_start..line_end];
        self.buffer_pos = line_end + 1;

        // Check for "data: [DONE]"
        if (std.mem.indexOf(u8, line, "[DONE]")) |_| {
            self.done = true;
            return StreamChunk{ .content = null, .done = true };
        }

        // Parse "data: {...}" line
        if (std.mem.startsWith(u8, line, "data: ")) {
            const json_str = line[6..];

            // Extract content field from JSON (simple parser)
            if (std.mem.indexOf(u8, json_str, "\"content\":\"")) |idx| {
                const content_start = idx + 11;
                if (std.mem.indexOf(u8, json_str[content_start..], "\"")) |content_end| {
                    const content = try self.allocator.dupe(u8, json_str[content_start .. content_start + content_end]);
                    return StreamChunk{ .content = content, .done = false };
                }
            }
        }

        // Skip empty lines or other data
        return try self.next();
    }

    pub fn deinit(self: *ChatStream) void {
        if (self.req) |*req| {
            req.deinit();
        }
        if (self.client) |*client| {
            client.deinit();
        }
    }
};

fn getEnv(key: []const u8) ?[]const u8 {
    return std.posix.getenv(key);
}

fn makeHTTPRequest(allocator: std.mem.Allocator, url: []const u8, json_payload: []const u8) ![]u8 {
    const uri = try std.Uri.parse(url);

    var client = http.Client{ .allocator = allocator };
    defer client.deinit();

    const server_header_buffer = try allocator.alloc(u8, 1024 * 8);
    defer allocator.free(server_header_buffer);

    var req = try client.open(.POST, uri, .{
        .server_header_buffer = server_header_buffer,
        .extra_headers = &[_]http.Header{
            .{ .name = "Content-Type", .value = "application/json" },
        },
    });
    defer req.deinit();

    req.transfer_encoding = .chunked;

    try req.send();
    try req.writeAll(json_payload);
    try req.finish();
    try req.wait();

    const body = try req.reader().readAllAlloc(allocator, 1024 * 1024 * 10);
    return body;
}

fn makeGETRequest(allocator: std.mem.Allocator, url: []const u8) ![]u8 {
    const uri = try std.Uri.parse(url);

    var client = http.Client{ .allocator = allocator };
    defer client.deinit();

    const server_header_buffer = try allocator.alloc(u8, 1024 * 8);
    defer allocator.free(server_header_buffer);

    var req = try client.open(.GET, uri, .{
        .server_header_buffer = server_header_buffer,
    });
    defer req.deinit();

    try req.send();
    try req.finish();
    try req.wait();

    const body = try req.reader().readAllAlloc(allocator, 1024 * 1024 * 10);
    return body;
}

// Helper function to extract content from non-streaming response
fn extractContent(response: []const u8) ?[]const u8 {
    if (std.mem.indexOf(u8, response, "\"content\":\"")) |idx| {
        const start = idx + 11;
        if (std.mem.indexOf(u8, response[start..], "\"")) |end| {
            return response[start .. start + end];
        }
    }
    return null;
}

pub fn main() !void {
    var gpa = std.heap.GeneralPurposeAllocator(.{}){};
    defer _ = gpa.deinit();
    const allocator = gpa.allocator();

    std.debug.print("UncloseAI Zig SDK - Production Library Example\n", .{});
    std.debug.print("===============================================\n\n", .{});

    // Initialize client with model discovery
    var client = try Client.init(allocator);
    defer client.deinit();

    const models = client.getModels();
    if (models.len == 0) {
        std.debug.print("\nERROR: No models discovered. Set environment variables:\n", .{});
        std.debug.print("  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.\n", .{});
        return error.NoModelsFound;
    }

    // Example 1: Non-streaming chat completion
    std.debug.print("\n=== Example 1: Non-Streaming Chat ===\n", .{});
    std.debug.print("Model: {s}\n", .{models[0].id});
    {
        const messages = [_]ChatMessage{
            .{ .role = "system", .content = "You are a helpful AI assistant." },
            .{ .role = "user", .content = "Explain quantum computing in one sentence." },
        };

        const request = ChatRequest{
            .model = models[0].id,
            .messages = &messages,
            .max_tokens = 100,
        };

        const response = client.chatCompletion(request) catch |err| {
            std.debug.print("Chat failed: {}\n", .{err});
            return err;
        };
        defer allocator.free(response);

        if (extractContent(response)) |content| {
            std.debug.print("Response: {s}\n", .{content});
        } else {
            std.debug.print("Response (raw): {s}\n", .{response});
        }
    }

    // Example 2: Streaming chat completion
    std.debug.print("\n=== Example 2: Streaming Chat ===\n", .{});
    std.debug.print("Model: {s}\n", .{models[0].id});
    {
        const messages = [_]ChatMessage{
            .{ .role = "system", .content = "You are a coding assistant." },
            .{ .role = "user", .content = "Write a hello world function in Zig." },
        };

        const request = ChatRequest{
            .model = models[0].id,
            .messages = &messages,
            .max_tokens = 200,
        };

        var stream = client.chatCompletionStream(request) catch |err| {
            std.debug.print("Stream creation failed: {}\n", .{err});
            return err;
        };
        defer stream.deinit();

        try stream.start();

        std.debug.print("Response (streaming): ", .{});
        while (try stream.next()) |chunk| {
            if (chunk.done) break;
            if (chunk.content) |content| {
                std.debug.print("{s}", .{content});
                var chunk_mut = chunk;
                chunk_mut.deinit(allocator);
            }
        }
        std.debug.print("\n", .{});
    }

    // Example 3: TTS Speech Generation
    const tts_endpoints = client.getTTSEndpoints();
    if (tts_endpoints.len > 0) {
        std.debug.print("\n=== Example 3: TTS Speech Generation ===\n", .{});
        std.debug.print("Endpoint: {s}\n", .{tts_endpoints[0]});

        const tts_request = TTSRequest{
            .model = "tts-1",
            .voice = "alloy",
            .input = "Hello from the UncloseAI Zig SDK! This library supports streaming completions and text to speech.",
        };

        const audio_data = client.generateSpeech(tts_request, 0) catch |err| {
            std.debug.print("TTS failed: {}\n", .{err});
            return err;
        };
        defer allocator.free(audio_data);

        const file = try std.fs.cwd().createFile("output.mp3", .{});
        defer file.close();
        try file.writeAll(audio_data);

        std.debug.print("Audio saved to output.mp3 ({} bytes)\n", .{audio_data.len});
    } else {
        std.debug.print("\n=== Example 3: TTS Speech Generation ===\n", .{});
        std.debug.print("No TTS endpoints available. Set TTS_ENDPOINT_1\n", .{});
    }

    std.debug.print("\n✅ UncloseAI Zig SDK examples complete!\n", .{});
}
