package llm

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

type Client struct {
	baseURL string
	apiKey  string
	model   string
	client  *http.Client
}

// Config holds optional per-request overrides for the LLM client.
// Empty fields are ignored, falling back to the client's defaults.
type Config struct {
	BaseURL string
	Model   string
	APIKey  string
}

type ChatMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type ChatRequest struct {
	Model              string        `json:"model"`
	Messages           []ChatMessage `json:"messages"`
	MaxCompletionTokens int          `json:"max_completion_tokens,omitempty"`
	Temperature        float64       `json:"temperature,omitempty"`
	Stream             bool          `json:"stream"`
}

type ChatResponse struct {
	ID      string `json:"id"`
	Object  string `json:"object"`
	Created int64  `json:"created"`
	Model   string `json:"model"`
	Choices []struct {
		Index        int         `json:"index"`
		Message      ChatMessage `json:"message"`
		FinishReason string      `json:"finish_reason"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}

type SSEEvent struct {
	Event string `json:"-"`
	Data  string `json:"data"`
}

func New(baseURL, apiKey, model string) *Client {
	return &Client{
		baseURL: strings.TrimRight(baseURL, "/") + "/v1",
		apiKey:  apiKey,
		model:   model,
		client: &http.Client{
			Timeout: 5 * time.Minute,
		},
	}
}

func (c *Client) Chat(messages []ChatMessage, maxTokens int, temperature float64, cfg *Config) (*ChatResponse, error) {
	model := c.model
	if cfg != nil && cfg.Model != "" {
		model = cfg.Model
	}
	req := ChatRequest{
		Model:              model,
		Messages:           messages,
		MaxCompletionTokens: maxTokens,
		Temperature:        temperature,
		Stream:             false,
	}

	return c.doChat(req, cfg)
}

func (c *Client) doChat(req ChatRequest, cfg *Config) (*ChatResponse, error) {
	body, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	baseURL := c.baseURL
	if cfg != nil && cfg.BaseURL != "" {
		baseURL = strings.TrimRight(cfg.BaseURL, "/") + "/v1"
	}

	httpReq, err := http.NewRequest("POST", baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return nil, fmt.Errorf("failed to create request: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	apiKey := c.apiKey
	if cfg != nil && cfg.APIKey != "" {
		apiKey = cfg.APIKey
	}
	httpReq.Header.Set("Authorization", "Bearer "+apiKey)

	resp, err := c.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	respBody, _ := io.ReadAll(resp.Body)

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API error (status %d): %s", resp.StatusCode, string(respBody))
	}

	var chatResp ChatResponse
	if err := json.Unmarshal(respBody, &chatResp); err != nil {
		return nil, fmt.Errorf("failed to parse response: %w", err)
	}

	return &chatResp, nil
}

func (c *Client) StreamChat(messages []ChatMessage, maxTokens int, temperature float64, callback func(string), cfg *Config) error {
	model := c.model
	if cfg != nil && cfg.Model != "" {
		model = cfg.Model
	}
	req := ChatRequest{
		Model:              model,
		Messages:           messages,
		MaxCompletionTokens: maxTokens,
		Temperature:        temperature,
		Stream:             true,
	}

	body, err := json.Marshal(req)
	if err != nil {
		return fmt.Errorf("failed to marshal request: %w", err)
	}

	baseURL := c.baseURL
	if cfg != nil && cfg.BaseURL != "" {
		baseURL = strings.TrimRight(cfg.BaseURL, "/") + "/v1"
	}

	httpReq, err := http.NewRequest("POST", baseURL+"/chat/completions", bytes.NewReader(body))
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}

	httpReq.Header.Set("Content-Type", "application/json")
	apiKey := c.apiKey
	if cfg != nil && cfg.APIKey != "" {
		apiKey = cfg.APIKey
	}
	httpReq.Header.Set("Authorization", "Bearer "+apiKey)

	resp, err := c.client.Do(httpReq)
	if err != nil {
		return fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

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

	return c.parseStream(resp.Body, callback)
}

func (c *Client) parseStream(body io.Reader, callback func(string)) error {
	scanner := bufio.NewScanner(body)
	for scanner.Scan() {
		line := strings.TrimSpace(scanner.Text())

		if line == "" || line == "data: [DONE]" {
			continue
		}

		if !strings.HasPrefix(line, "data: ") {
			continue
		}

		data := strings.TrimPrefix(line, "data: ")
		var choice struct {
			Delta struct {
				Content string `json:"content"`
			} `json:"delta"`
		}
		if err := json.Unmarshal([]byte(data), &choice); err != nil {
			continue
		}

		if choice.Delta.Content != "" {
			callback(choice.Delta.Content)
		}
	}

	return scanner.Err()
}