package ai

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

// openAICompatible 是 OpenAI Chat Completions 兼容端点的通用流式客户端；
// 讯飞星火 Lite 与 DeepSeek 均暴露该协议，只是 baseURL/model/key 不同。
type openAICompatible struct {
	name    string
	baseURL string
	model   string
	apiKey  string
	client  *http.Client
}

func (p *openAICompatible) Name() string { return p.name }

type chatRequest struct {
	Model         string         `json:"model"`
	Messages      []Message      `json:"messages"`
	Stream        bool           `json:"stream"`
	StreamOptions *streamOptions `json:"stream_options,omitempty"`
}

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

type chatDelta struct {
	Choices []struct {
		Delta struct {
			Content string `json:"content"`
		} `json:"delta"`
	} `json:"choices"`
	Usage *chatUsage `json:"usage"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error"`
}

// chatUsage 兼容 OpenAI / DeepSeek / 星火：缓存命中字段名不统一。
type chatUsage struct {
	PromptTokens         int64 `json:"prompt_tokens"`
	CompletionTokens     int64 `json:"completion_tokens"`
	PromptCacheHitTokens int64 `json:"prompt_cache_hit_tokens"`
	PromptTokensDetails  *struct {
		CachedTokens int64 `json:"cached_tokens"`
	} `json:"prompt_tokens_details"`
}

func parseUsage(u *chatUsage) *Usage {
	if u == nil {
		return nil
	}
	cached := u.PromptCacheHitTokens
	if u.PromptTokensDetails != nil && u.PromptTokensDetails.CachedTokens > cached {
		cached = u.PromptTokensDetails.CachedTokens
	}
	return &Usage{
		PromptTokens:     u.PromptTokens,
		CompletionTokens: u.CompletionTokens,
		CachedTokens:     cached,
	}
}

func (p *openAICompatible) ChatStream(ctx context.Context, messages []Message) (<-chan Chunk, error) {
	body, e := json.Marshal(chatRequest{
		Model: p.model, Messages: messages, Stream: true,
		StreamOptions: &streamOptions{IncludeUsage: true},
	})
	if e != nil {
		return nil, e
	}
	req, e := http.NewRequestWithContext(ctx, http.MethodPost, strings.TrimRight(p.baseURL, "/")+"/chat/completions", bytes.NewReader(body))
	if e != nil {
		return nil, e
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+p.apiKey)
	req.Header.Set("Accept", "text/event-stream")
	resp, e := p.client.Do(req)
	if e != nil {
		return nil, errors.New("AI_NETWORK_ERROR")
	}
	if resp.StatusCode != http.StatusOK {
		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
		resp.Body.Close()
		if resp.StatusCode == http.StatusUnauthorized || resp.StatusCode == http.StatusForbidden {
			return nil, errors.New("AI_KEY_INVALID")
		}
		return nil, fmt.Errorf("AI_HTTP_%d: %s", resp.StatusCode, truncate(string(raw), 300))
	}

	out := make(chan Chunk, 16)
	go func() {
		defer close(out)
		defer resp.Body.Close()
		sc := bufio.NewScanner(resp.Body)
		sc.Buffer(make([]byte, 0, 64*1024), 1024*1024)
		var promptChars, completionChars int
		for _, m := range messages {
			promptChars += len(m.Content)
		}
		var lastUsage *Usage
		for sc.Scan() {
			line := strings.TrimSpace(sc.Text())
			if !strings.HasPrefix(line, "data:") {
				continue
			}
			payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
			if payload == "" || payload == "[DONE]" {
				continue
			}
			var d chatDelta
			if json.Unmarshal([]byte(payload), &d) != nil {
				continue
			}
			if d.Error != nil {
				out <- Chunk{Err: errors.New(truncate(d.Error.Message, 300))}
				return
			}
			if d.Usage != nil {
				lastUsage = parseUsage(d.Usage)
			}
			if len(d.Choices) > 0 && d.Choices[0].Delta.Content != "" {
				completionChars += len(d.Choices[0].Delta.Content)
				select {
				case out <- Chunk{Content: d.Choices[0].Delta.Content}:
				case <-ctx.Done():
					return
				}
			}
		}
		if e := sc.Err(); e != nil && ctx.Err() == nil {
			out <- Chunk{Err: errors.New("AI_STREAM_INTERRUPTED")}
			return
		}
		if lastUsage == nil {
			lastUsage = &Usage{
				PromptTokens:     int64((promptChars + 3) / 4),
				CompletionTokens: int64((completionChars + 3) / 4),
				Estimated:        true,
			}
		}
		select {
		case out <- Chunk{Usage: lastUsage}:
		case <-ctx.Done():
		}
	}()
	return out, nil
}

func truncate(s string, n int) string {
	if len(s) <= n {
		return s
	}
	return s[:n] + "..."
}

func newHTTPClient() *http.Client {
	return &http.Client{Timeout: 5 * time.Minute}
}