// This is free software for the public good of a permacomputer hosted at
// permacomputer.com, an always-on computer by the people, for the people.
// One which is durable, easy to repair, & distributed like tap water
// for machine learning intelligence.
//
// The permacomputer is community-owned infrastructure optimized around
// four values:
//
//   TRUTH      First principles, math & science, open source code freely distributed
//   FREEDOM    Voluntary partnerships, freedom from tyranny & corporate control
//   HARMONY    Minimal waste, self-renewing systems with diverse thriving connections
//   LOVE       Be yourself without hurting others, cooperation through natural law
//
// This software contributes to that vision by making machine learning
// accessible to everyone through a free, open, embeddable chat interface.
// Code is seeds to sprout on any abandoned technology.

package uncloseai

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

// Error types
type Error string

const (
	ErrConnection      Error = "connection error"
	ErrModelNotFound   Error = "model not found"
	ErrStreaming       Error = "streaming error"
	ErrNoModels        Error = "no models available"
	ErrNoTTSEndpoints  Error = "no TTS endpoints available"
	ErrInvalidResponse Error = "invalid response"
)

func (e Error) Error() string {
	return string(e)
}

// ModelInfo contains metadata about a discovered model
type ModelInfo struct {
	ID        string `json:"id"`
	Endpoint  string `json:"endpoint"`
	MaxTokens int    `json:"max_tokens"`
}

// Message represents a chat message
type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// ChatResponse is the response from a non-streaming chat completion
type ChatResponse struct {
	ID      string       `json:"id"`
	Model   string       `json:"model"`
	Choices []ChatChoice `json:"choices"`
}

// ChatChoice represents a single choice in the response
type ChatChoice struct {
	Index        int     `json:"index"`
	Message      Message `json:"message"`
	FinishReason string  `json:"finish_reason"`
}

// StreamChunk represents a chunk from streaming chat
type StreamChunk struct {
	ID      string             `json:"id"`
	Model   string             `json:"model"`
	Choices []StreamChunkChoice `json:"choices"`
}

// StreamChunkChoice represents a single choice in a streaming chunk
type StreamChunkChoice struct {
	Index int         `json:"index"`
	Delta StreamDelta `json:"delta"`
}

// StreamDelta represents the delta content in a streaming chunk
type StreamDelta struct {
	Role    string `json:"role,omitempty"`
	Content string `json:"content,omitempty"`
}

// Config holds configuration options for the UncloseAI client
type Config struct {
	// Endpoints for model discovery (nil = auto-discover from environment)
	Endpoints []string
	// TTS endpoints (nil = auto-discover from environment)
	TTSEndpoints []string
	// API key for authentication (empty = no authentication)
	APIKey string
	// HTTP client timeout (0 = 30 seconds default)
	Timeout time.Duration
	// Enable debug logging
	Debug bool
}

// ChatOptions holds options for chat completions
type ChatOptions struct {
	// Model ID (empty = auto-select first available)
	Model string
	// Maximum tokens to generate (0 = no limit)
	MaxTokens int
	// Sampling temperature (0.0 - 2.0)
	Temperature float64
	// Nucleus sampling (0.0 - 1.0)
	TopP float64
}

// Client is the main UncloseAI client
type Client struct {
	models       []ModelInfo
	ttsEndpoints []string
	apiKey       string
	timeout      time.Duration
	httpClient   *http.Client
	debug        bool
}

// New creates a new UncloseAI client with optional configuration.
// If config is nil, defaults are used and endpoints are auto-discovered from environment variables.
func New(config *Config) (*Client, error) {
	if config == nil {
		config = &Config{}
	}

	timeout := config.Timeout
	if timeout == 0 {
		timeout = 30 * time.Second
	}

	client := &Client{
		models:       make([]ModelInfo, 0),
		ttsEndpoints: make([]string, 0),
		apiKey:       config.APIKey,
		timeout:      timeout,
		httpClient:   &http.Client{Timeout: timeout},
		debug:        config.Debug,
	}

	// Discover endpoints from environment if not provided
	endpoints := config.Endpoints
	if endpoints == nil {
		endpoints = discoverEnvEndpoints("MODEL_ENDPOINT")
	}

	ttsEndpoints := config.TTSEndpoints
	if ttsEndpoints == nil {
		ttsEndpoints = discoverEnvEndpoints("TTS_ENDPOINT")
	}

	if client.debug {
		fmt.Printf("[DEBUG] Initialized with %d model endpoint(s) and %d TTS endpoint(s)\n",
			len(endpoints), len(ttsEndpoints))
	}

	// Discover models from each endpoint
	for _, endpoint := range endpoints {
		if err := client.discoverModelsFromEndpoint(endpoint); err != nil {
			if client.debug {
				fmt.Printf("[DEBUG] Failed to discover models from %s: %v\n", endpoint, err)
			}
		}
	}

	client.ttsEndpoints = ttsEndpoints

	return client, nil
}

// discoverEnvEndpoints finds endpoints from environment variables like PREFIX_1, PREFIX_2, ...
func discoverEnvEndpoints(prefix string) []string {
	endpoints := make([]string, 0)
	for i := 1; i <= 9999; i++ {
		endpoint := os.Getenv(fmt.Sprintf("%s_%d", prefix, i))
		if endpoint == "" {
			break
		}
		endpoints = append(endpoints, endpoint)
	}
	return endpoints
}

// discoverModelsFromEndpoint discovers available models from an endpoint
func (c *Client) discoverModelsFromEndpoint(endpoint string) error {
	if c.debug {
		fmt.Printf("[DEBUG] Discovering models from: %s\n", endpoint)
	}

	req, err := http.NewRequest("GET", endpoint+"/models", nil)
	if err != nil {
		return err
	}

	if c.apiKey != "" {
		req.Header.Set("Authorization", "Bearer "+c.apiKey)
	}

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("unexpected status: %d", resp.StatusCode)
	}

	var result struct {
		Data []struct {
			ID          string `json:"id"`
			MaxModelLen int    `json:"max_model_len"`
		} `json:"data"`
	}

	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return err
	}

	for _, model := range result.Data {
		maxTokens := model.MaxModelLen
		if maxTokens == 0 {
			maxTokens = 8192
		}
		c.models = append(c.models, ModelInfo{
			ID:        model.ID,
			Endpoint:  endpoint,
			MaxTokens: maxTokens,
		})

		if c.debug {
			fmt.Printf("[DEBUG] Discovered: %s\n", model.ID)
		}
	}

	return nil
}

// ListModels returns all discovered models
func (c *Client) ListModels() []ModelInfo {
	return c.models
}

// Chat performs a non-streaming chat completion
func (c *Client) Chat(ctx context.Context, messages []Message, options *ChatOptions) (*ChatResponse, error) {
	if options == nil {
		options = &ChatOptions{Temperature: 0.7, TopP: 1.0}
	}

	modelInfo, err := c.getModelInfo(options.Model)
	if err != nil {
		return nil, err
	}

	payload := map[string]interface{}{
		"model":       modelInfo.ID,
		"messages":    messages,
		"temperature": options.Temperature,
		"top_p":       options.TopP,
		"stream":      false,
	}

	if options.MaxTokens > 0 {
		payload["max_tokens"] = options.MaxTokens
	}

	body, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, "POST", modelInfo.Endpoint+"/chat/completions", bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}

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

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrConnection, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		bodyBytes, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("%w: status %d: %s", ErrInvalidResponse, resp.StatusCode, string(bodyBytes))
	}

	var chatResp ChatResponse
	if err := json.NewDecoder(resp.Body).Decode(&chatResp); err != nil {
		return nil, fmt.Errorf("%w: %v", ErrInvalidResponse, err)
	}

	return &chatResp, nil
}

// ChatStream performs a streaming chat completion, returning channels for chunks and errors
func (c *Client) ChatStream(ctx context.Context, messages []Message, options *ChatOptions) (<-chan StreamChunk, <-chan error) {
	chunkChan := make(chan StreamChunk)
	errChan := make(chan error, 1)

	go func() {
		defer close(chunkChan)
		defer close(errChan)

		if options == nil {
			options = &ChatOptions{Temperature: 0.7, TopP: 1.0}
		}

		modelInfo, err := c.getModelInfo(options.Model)
		if err != nil {
			errChan <- err
			return
		}

		payload := map[string]interface{}{
			"model":       modelInfo.ID,
			"messages":    messages,
			"temperature": options.Temperature,
			"top_p":       options.TopP,
			"stream":      true,
		}

		if options.MaxTokens > 0 {
			payload["max_tokens"] = options.MaxTokens
		}

		body, err := json.Marshal(payload)
		if err != nil {
			errChan <- err
			return
		}

		req, err := http.NewRequestWithContext(ctx, "POST", modelInfo.Endpoint+"/chat/completions", bytes.NewBuffer(body))
		if err != nil {
			errChan <- err
			return
		}

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

		resp, err := c.httpClient.Do(req)
		if err != nil {
			errChan <- fmt.Errorf("%w: %v", ErrConnection, err)
			return
		}
		defer resp.Body.Close()

		if resp.StatusCode != http.StatusOK {
			bodyBytes, _ := io.ReadAll(resp.Body)
			errChan <- fmt.Errorf("%w: status %d: %s", ErrStreaming, resp.StatusCode, string(bodyBytes))
			return
		}

		// Parse SSE stream
		scanner := bufio.NewScanner(resp.Body)
		for scanner.Scan() {
			line := scanner.Text()

			// SSE format: "data: {...}"
			if strings.HasPrefix(line, "data: ") {
				data := strings.TrimPrefix(line, "data: ")

				// Check for stream termination
				if strings.TrimSpace(data) == "[DONE]" {
					break
				}

				var chunk StreamChunk
				if err := json.Unmarshal([]byte(data), &chunk); err != nil {
					if c.debug {
						fmt.Printf("[DEBUG] Failed to parse chunk: %s\n", data)
					}
					continue // Skip malformed chunks
				}

				select {
				case chunkChan <- chunk:
				case <-ctx.Done():
					return
				}
			}
		}

		if err := scanner.Err(); err != nil {
			errChan <- fmt.Errorf("%w: %v", ErrStreaming, err)
		}
	}()

	return chunkChan, errChan
}

// TTS generates speech from text
func (c *Client) TTS(ctx context.Context, text, voice, model string) ([]byte, error) {
	if len(c.ttsEndpoints) == 0 {
		return nil, ErrNoTTSEndpoints
	}

	endpoint := c.ttsEndpoints[0]

	payload := map[string]interface{}{
		"model": model,
		"voice": voice,
		"input": text,
	}

	body, err := json.Marshal(payload)
	if err != nil {
		return nil, err
	}

	req, err := http.NewRequestWithContext(ctx, "POST", endpoint+"/audio/speech", bytes.NewBuffer(body))
	if err != nil {
		return nil, err
	}

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

	resp, err := c.httpClient.Do(req)
	if err != nil {
		return nil, fmt.Errorf("%w: %v", ErrConnection, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("%w: status %d", ErrInvalidResponse, resp.StatusCode)
	}

	return io.ReadAll(resp.Body)
}

// getModelInfo retrieves model info by ID or returns first available
func (c *Client) getModelInfo(modelID string) (*ModelInfo, error) {
	if len(c.models) == 0 {
		return nil, ErrNoModels
	}

	if modelID == "" {
		return &c.models[0], nil
	}

	for i := range c.models {
		if c.models[i].ID == modelID {
			return &c.models[i], nil
		}
	}

	return nil, fmt.Errorf("%w: '%s'", ErrModelNotFound, modelID)
}