package llm

import (
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"net/http"
	"strings"

	"ai-operator/internal/config"
)

type OpenAIStreaming struct {
	cfg    config.Config
	client *http.Client
}

func NewOpenAIStreaming(cfg config.Config) *OpenAIStreaming {
	return &OpenAIStreaming{cfg: cfg, client: &http.Client{Timeout: cfg.LLM.Timeout}}
}

func (p *OpenAIStreaming) StreamGenerate(ctx context.Context, req GenerateRequest) (<-chan Event, error) {
	if p.cfg.OpenAI.APIKey == "" {
		return nil, errors.New("OPENAI_API_KEY is required for streaming LLM")
	}
	out := make(chan Event, 64)
	go p.run(ctx, req, out)
	return out, nil
}

func (p *OpenAIStreaming) run(ctx context.Context, req GenerateRequest, out chan<- Event) {
	defer close(out)
	payload := map[string]any{
		"model":       p.cfg.LLM.Model,
		"temperature": p.cfg.LLM.Temperature,
		"max_tokens":  p.cfg.LLM.MaxOutputTokens,
		"stream":      true,
	}
	msgs := []map[string]string{{"role": "system", "content": req.SystemPrompt}}
	for _, m := range req.Messages {
		if strings.TrimSpace(m.Content) != "" {
			msgs = append(msgs, map[string]string{"role": m.Role, "content": m.Content})
		}
	}
	payload["messages"] = msgs
	body, _ := json.Marshal(payload)
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, "https://api.openai.com/v1/chat/completions", bytes.NewReader(body))
	if err != nil {
		out <- Event{Type: EventError, Error: err.Error()}
		return
	}
	httpReq.Header.Set("Authorization", "Bearer "+p.cfg.OpenAI.APIKey)
	httpReq.Header.Set("Content-Type", "application/json")
	resp, err := p.client.Do(httpReq)
	if err != nil {
		out <- Event{Type: EventError, Error: err.Error()}
		return
	}
	defer resp.Body.Close()
	if resp.StatusCode < 200 || resp.StatusCode > 299 {
		out <- Event{Type: EventError, Error: "openai streaming llm returned non-2xx"}
		return
	}
	sc := bufio.NewScanner(resp.Body)
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if !strings.HasPrefix(line, "data:") {
			continue
		}
		data := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
		if data == "[DONE]" {
			out <- Event{Type: EventDone}
			return
		}
		var chunk struct {
			Choices []struct {
				Delta struct {
					Content string `json:"content"`
				} `json:"delta"`
			} `json:"choices"`
		}
		if err := json.Unmarshal([]byte(data), &chunk); err != nil {
			continue
		}
		for _, c := range chunk.Choices {
			if c.Delta.Content != "" {
				out <- Event{Type: EventTextDelta, Text: c.Delta.Content}
			}
		}
	}
	if err := sc.Err(); err != nil {
		out <- Event{Type: EventError, Error: err.Error()}
		return
	}
	out <- Event{Type: EventDone}
}