package llm

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

	"laodingbot/internal/config"
	"laodingbot/internal/logger"
)

type Client interface {
	Generate(ctx context.Context, systemPrompt, userPrompt string) (string, error)
}

type OpenAICompatibleClient struct {
	baseURL string
	apiKey  string
	model   string
	http    *http.Client
	log     *logger.Logger
}

func NewOpenAICompatibleClient(cfg config.LLMConfig, log *logger.Logger) *OpenAICompatibleClient {
	return &OpenAICompatibleClient{
		baseURL: cfg.BaseURL,
		apiKey:  cfg.APIKey,
		model:   cfg.Model,
		http:    &http.Client{Timeout: 60 * time.Second},
		log:     log,
	}
}

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
}

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

type chatResponse struct {
	Choices []struct {
		Message chatMessage `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

func (c *OpenAICompatibleClient) Generate(ctx context.Context, systemPrompt, userPrompt string) (string, error) {
	if c.log != nil {
		c.log.Debugf("llm request start model=%s system_len=%d user_len=%d", c.model, len(systemPrompt), len(userPrompt))
	}
	body := chatRequest{
		Model: c.model,
		Messages: []chatMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userPrompt},
		},
	}
	b, err := json.Marshal(body)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("marshal llm request failed err=%v", err)
		}
		return "", err
	}

	url := strings.TrimRight(c.baseURL, "/") + "/chat/completions"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		if c.log != nil {
			c.log.Errorf("build llm request failed err=%v", err)
		}
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.http.Do(req)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("llm http request failed err=%v", err)
		}
		return "", err
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("llm read response failed err=%v", err)
		}
		return "", err
	}

	var out chatResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		if c.log != nil {
			c.log.Errorf("llm response unmarshal failed status=%d err=%v", resp.StatusCode, err)
		}
		return "", err
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if c.log != nil {
			c.log.Errorf("llm bad status=%d", resp.StatusCode)
		}
		if out.Error != nil && out.Error.Message != "" {
			return "", fmt.Errorf("llm error: %s", out.Error.Message)
		}
		return "", fmt.Errorf("llm error status: %d", resp.StatusCode)
	}

	if len(out.Choices) == 0 {
		if c.log != nil {
			c.log.Errorf("llm returned empty choices status=%d", resp.StatusCode)
		}
		return "", fmt.Errorf("llm returned empty choices")
	}
	if c.log != nil {
		c.log.Infof("llm response success model=%s output_len=%d", c.model, len(out.Choices[0].Message.Content))
	}

	return out.Choices[0].Message.Content, nil
}