package llm

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

// openAI 是 OpenAI 兼容协议（/chat/completions + SSE）的实现。
type openAI struct {
	opts Options
	http *http.Client
}

func newOpenAI(opts Options) *openAI {
	if opts.MaxTokens <= 0 {
		opts.MaxTokens = 1024
	}
	if opts.Temperature <= 0 {
		opts.Temperature = 0.3
	}
	return &openAI{
		opts: opts,
		http: &http.Client{
			// 不设总超时：流式响应可能持续很久，由 ctx 控制。
			Transport: &http.Transport{
				Proxy:               http.ProxyFromEnvironment,
				MaxIdleConns:        8,
				IdleConnTimeout:     90 * time.Second,
				ForceAttemptHTTP2:   true,
				TLSHandshakeTimeout: 15 * time.Second,
			},
		},
	}
}

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

type contentPart struct {
	Type     string    `json:"type"`
	Text     string    `json:"text,omitempty"`
	ImageURL *imageURL `json:"image_url,omitempty"`
}

type imageURL struct {
	URL string `json:"url"`
}

type streamOptions struct {
	IncludeUsage bool `json:"include_usage"`
}

type chatRequest struct {
	Model         string         `json:"model"`
	Messages      []reqMessage   `json:"messages"`
	Stream        bool           `json:"stream"`
	MaxTokens     int            `json:"max_tokens,omitempty"`
	Temperature   float64        `json:"temperature,omitempty"`
	StreamOptions *streamOptions `json:"stream_options,omitempty"`
}

type streamChunk struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
		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"`
	Error *struct {
		Message string `json:"message"`
		Type    string `json:"type"`
	} `json:"error"`
}

func (o *openAI) endpoint() string {
	base := strings.TrimRight(o.opts.BaseURL, "/")
	if strings.HasSuffix(base, "/chat/completions") {
		return base
	}
	return base + "/chat/completions"
}

func (o *openAI) setHeaders(req *http.Request, stream bool) {
	req.Header.Set("Content-Type", "application/json")
	if stream {
		req.Header.Set("Accept", "text/event-stream")
	} else {
		req.Header.Set("Accept", "application/json")
	}
	if o.opts.APIKey != "" {
		req.Header.Set("Authorization", "Bearer "+o.opts.APIKey)
	}
	for k, v := range o.opts.ExtraHeaders {
		req.Header.Set(k, v)
	}
}

func (o *openAI) buildBody(system string, msgs []Message, stream bool) ([]byte, error) {
	body := chatRequest{
		Model:       o.opts.Model,
		Stream:      stream,
		MaxTokens:   o.opts.MaxTokens,
		Temperature: o.opts.Temperature,
	}
	if stream {
		body.StreamOptions = &streamOptions{IncludeUsage: true}
	}
	if strings.TrimSpace(system) != "" {
		body.Messages = append(body.Messages, reqMessage{Role: "system", Content: system})
	}
	for _, m := range msgs {
		role := string(m.Role)
		if role == "" {
			role = "user"
		}
		if len(m.Image) == 0 {
			body.Messages = append(body.Messages, reqMessage{Role: role, Content: m.Text})
			continue
		}
		text := m.Text
		if strings.TrimSpace(text) == "" {
			text = "请看图并回答。"
		}
		parts := []contentPart{
			{Type: "text", Text: text},
			{
				Type: "image_url",
				ImageURL: &imageURL{
					URL: "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(m.Image),
				},
			},
		}
		body.Messages = append(body.Messages, reqMessage{Role: role, Content: parts})
	}
	raw, err := json.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("构造请求失败: %w", err)
	}
	return raw, nil
}

// Chat 发起一次流式对话。
func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error) {
	if len(msgs) == 0 {
		return Usage{}, fmt.Errorf("没有可发送的消息")
	}
	payload, err := o.buildBody(system, msgs, true)
	if err != nil {
		return Usage{}, err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
	if err != nil {
		return Usage{}, fmt.Errorf("创建请求失败: %w", err)
	}
	o.setHeaders(req, true)

	resp, err := o.http.Do(req)
	if err != nil {
		return Usage{}, fmt.Errorf("请求 %s 失败: %w", o.opts.Name, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
		return Usage{}, fmt.Errorf("模型返回 HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}

	// 部分服务端会忽略 stream 参数直接返回完整 JSON，这里做兼容。
	if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "event-stream") {
		return o.readWhole(resp.Body, fn)
	}

	var usage Usage
	var got strings.Builder
	err = readSSE(resp.Body, func(ev sseEvent) error {
		if len(ev.Data) == 0 || bytes.Equal(ev.Data, []byte("[DONE]")) {
			return nil
		}
		var chunk streamChunk
		if json.Unmarshal(ev.Data, &chunk) != nil {
			return nil // 忽略无法解析的心跳/杂项
		}
		if chunk.Error != nil {
			return fmt.Errorf("模型返回错误: %s", chunk.Error.Message)
		}
		if chunk.Usage != nil {
			usage = Usage{
				PromptTokens:     chunk.Usage.PromptTokens,
				CompletionTokens: chunk.Usage.CompletionTokens,
				TotalTokens:      chunk.Usage.TotalTokens,
			}
		}
		for _, c := range chunk.Choices {
			if c.Delta.Content == "" {
				continue
			}
			got.WriteString(c.Delta.Content)
			if fn != nil {
				if derr := fn(c.Delta.Content); derr != nil {
					return derr
				}
			}
		}
		return nil
	})
	if err != nil {
		return usage, err
	}
	if got.Len() == 0 {
		return usage, fmt.Errorf("模型没有返回任何内容，请确认 %s 是支持图片输入的视觉模型", o.opts.Model)
	}
	return usage, nil
}

// readWhole 处理服务端忽略 stream 参数的情况。
func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
	raw, err := io.ReadAll(io.LimitReader(r, 8<<20))
	if err != nil {
		return Usage{}, fmt.Errorf("读取响应失败: %w", err)
	}
	var resp struct {
		Choices []struct {
			Message struct {
				Content string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage *struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
			TotalTokens      int `json:"total_tokens"`
		} `json:"usage"`
		Error *struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	if err := json.Unmarshal(raw, &resp); err != nil {
		return Usage{}, fmt.Errorf("解析响应失败: %w", err)
	}
	if resp.Error != nil {
		return Usage{}, fmt.Errorf("模型返回错误: %s", resp.Error.Message)
	}
	if len(resp.Choices) == 0 || resp.Choices[0].Message.Content == "" {
		return Usage{}, fmt.Errorf("模型返回内容为空")
	}
	var usage Usage
	if resp.Usage != nil {
		usage = Usage{
			PromptTokens:     resp.Usage.PromptTokens,
			CompletionTokens: resp.Usage.CompletionTokens,
			TotalTokens:      resp.Usage.TotalTokens,
		}
	}
	text := resp.Choices[0].Message.Content
	if fn != nil {
		if err := fn(text); err != nil {
			return usage, err
		}
	}
	return usage, nil
}

// Ping 做一次最小代价的连通性检查。
func (o *openAI) Ping(ctx context.Context) error {
	payload, err := o.buildBody("", []Message{{Role: RoleUser, Text: "ping"}}, false)
	if err != nil {
		return err
	}
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
	if err != nil {
		return fmt.Errorf("创建请求失败: %w", err)
	}
	o.setHeaders(req, false)

	resp, err := o.http.Do(req)
	if err != nil {
		return fmt.Errorf("连接 %s 失败: %w", o.opts.BaseURL, err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
	}
	if len(raw) == 0 {
		return fmt.Errorf("服务返回空响应")
	}
	return nil
}