package main

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"
)

// Configuration - Replace with your API credentials
const (
	// Option 1: OpenAI API
	apiKey    = "YOUR_API_KEY_HERE"
	apiURL    = "https://api.openai.com/v1/chat/completions"
	modelName = "gpt-3.5-turbo" // or gpt-4, etc.

	// Alternative: an existing OpenAI-compatible local server (for example,
	// llamafile at http://127.0.0.1:8080/v1/chat/completions, or Ollama at
	// http://127.0.0.1:11434/v1/chat/completions). Set apiKey to "" and select
	// a model already available on that server. This adapter does not speak
	// Anthropic's native /v1/messages or Ollama's native /api/chat protocol.
)

// LLM calls the language model. If stream is nil, returns complete response via return value.
// If stream is provided, streams response chunks to channel and returns empty string.
// Input can be a string (wrapped as user message) or []map[string]string for full message history.
func LLM(ctx context.Context, input interface{}, stream chan<- string) (string, error) {
	// The adapter owns the stream and must close it on every return path,
	// including invalid input, HTTP errors and cancellation.
	if stream != nil {
		defer close(stream)
	}
	// Build messages array
	var messages []map[string]string
	switch v := input.(type) {
	case string:
		messages = []map[string]string{
			{"role": "user", "content": v},
		}
	case []map[string]string:
		messages = v
	default:
		return "", fmt.Errorf("invalid input type")
	}

	// Build request
	requestBody := map[string]interface{}{
		"model":       modelName,
		"messages":    messages,
		"temperature": 0.7,
		"max_tokens":  500,
	}

	if stream != nil {
		requestBody["stream"] = true
	}

	jsonBody, err := json.Marshal(requestBody)
	if err != nil {
		return "", err
	}

	req, err := http.NewRequestWithContext(ctx, "POST", apiURL, bytes.NewBuffer(jsonBody))
	if err != nil {
		return "", err
	}

	req.Header.Set("Content-Type", "application/json")
	if apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+apiKey)
	}

	client := &http.Client{Timeout: 60 * time.Second}
	resp, err := client.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		return "", fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(body))
	}

	// Handle streaming response
	if stream != nil {
		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			line := scanner.Text()
			if strings.HasPrefix(line, "data: ") {
				data := strings.TrimPrefix(line, "data: ")
				if data == "[DONE]" {
					return "", nil
				}

				var chunk map[string]interface{}
				if err := json.Unmarshal([]byte(data), &chunk); err == nil {
					if choices, ok := chunk["choices"].([]interface{}); ok && len(choices) > 0 {
						if choice, ok := choices[0].(map[string]interface{}); ok {
							if delta, ok := choice["delta"].(map[string]interface{}); ok {
								if content, ok := delta["content"].(string); ok {
									select {
									case stream <- content:
									case <-ctx.Done():
										return "", ctx.Err()
									}
								}
							}
						}
					}
				}
			}
		}
		return "", scanner.Err()
	}

	// Handle non-streaming response
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	var response map[string]interface{}
	if err := json.Unmarshal(body, &response); err != nil {
		return "", err
	}

	if choices, ok := response["choices"].([]interface{}); ok && len(choices) > 0 {
		if choice, ok := choices[0].(map[string]interface{}); ok {
			if message, ok := choice["message"].(map[string]interface{}); ok {
				if content, ok := message["content"].(string); ok {
					return content, nil
				}
			}
		}
	}

	return "", fmt.Errorf("unexpected response format")
}