package backend

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

// OpenAIBackend speaks the OpenAI /v1/chat/completions wire format.
// Compatible with OpenRouter, OpenAI, and any other OpenAI-compatible API.
type OpenAIBackend struct {
	name         string
	baseURL      string
	apiKey       string
	model        string
	client       *http.Client
	extraHeaders map[string]string // optional; applied after Authorization
}

func NewOpenAI(name, baseURL, apiKey string) *OpenAIBackend {
	if baseURL == "" {
		baseURL = "https://api.openai.com"
	}
	b := &OpenAIBackend{name: name, baseURL: baseURL, apiKey: apiKey, client: &http.Client{}}
	b.model = b.DefaultModel()
	return b
}

func (b *OpenAIBackend) Name() string         { return b.name }
func (b *OpenAIBackend) Model() string        { return b.model }
func (b *OpenAIBackend) SetModel(m string)    { b.model = m }
func (b *OpenAIBackend) DefaultModel() string {
	switch b.name {
	case "openrouter":
		return "deepseek/deepseek-v3.2"
	case "anthropic":
		return "claude-sonnet-4-5"
	default:
		return "qwen3.5:9b"
	}
}

// -- wire types --

type openAIMessage struct {
	Role       string           `json:"role"`
	Content    *string          `json:"content"`
	ToolCalls  []openAIToolCall `json:"tool_calls,omitempty"`
	ToolCallID string           `json:"tool_call_id,omitempty"`
}

type openAIToolCall struct {
	ID       string             `json:"id"`
	Index    int                `json:"index,omitempty"`
	Type     string             `json:"type"`
	Function openAIFunctionCall `json:"function"`
}

type openAIFunctionCall struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"` // JSON string, not object
}

type openAITool struct {
	Type     string             `json:"type"`
	Function openAIToolFunction `json:"function"`
}

type openAIToolFunction struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters"`
}

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

type openAIChatRequest struct {
	Model            string               `json:"model"`
	Messages         []openAIMessage      `json:"messages"`
	Tools            []openAITool         `json:"tools,omitempty"`
	Stream           bool                 `json:"stream"`
	StreamOptions    *openAIStreamOptions `json:"stream_options,omitempty"`
	MaxTokens        int                  `json:"max_tokens,omitempty"`
	Temperature      *float64             `json:"temperature,omitempty"`
	FrequencyPenalty *float64             `json:"frequency_penalty,omitempty"`
	PresencePenalty  *float64             `json:"presence_penalty,omitempty"`
}

type openAIUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
}

type openAIDelta struct {
	Role      string           `json:"role"`
	Content   string           `json:"content"`
	ToolCalls []openAIToolCall `json:"tool_calls,omitempty"`
}

type openAIStreamChoice struct {
	Index        int         `json:"index"`
	Delta        openAIDelta `json:"delta"`
	FinishReason string      `json:"finish_reason"`
}

type openAIStreamResponse struct {
	Choices []openAIStreamChoice `json:"choices"`
	Usage   openAIUsage          `json:"usage"`
}

// -- implementation --

func (b *OpenAIBackend) ChatStream(ctx context.Context, messages []Message, tools []Tool, params GenerationParams) (<-chan StreamEvent, error) {
	model := b.model
	wireMessages := make([]openAIMessage, len(messages))
	for i, m := range messages {
		wm := openAIMessage{
			Role:       m.Role,
			ToolCallID: m.ToolCallID,
		}
		for _, tc := range m.ToolCalls {
			wm.ToolCalls = append(wm.ToolCalls, openAIToolCall{
				ID:   tc.ID,
				Type: "function",
				Function: openAIFunctionCall{
					Name:      tc.Name,
					Arguments: string(tc.Arguments),
				},
			})
		}
		// OpenAI spec: content must be null (not "") when tool_calls is present.
		if len(wm.ToolCalls) == 0 {
			wm.Content = &m.Content
		}
		wireMessages[i] = wm
	}

	var wireTools []openAITool
	for _, t := range tools {
		wireTools = append(wireTools, openAITool{
			Type: "function",
			Function: openAIToolFunction{
				Name:        t.Name,
				Description: t.Description,
				Parameters:  t.Parameters,
			},
		})
	}

	req := openAIChatRequest{
		Model:            model,
		Messages:         wireMessages,
		Tools:            wireTools,
		Stream:           true,
		StreamOptions:    &openAIStreamOptions{IncludeUsage: true},
		MaxTokens:        params.MaxTokens,
		Temperature:      params.Temperature,
		FrequencyPenalty: params.FrequencyPenalty,
		PresencePenalty:  params.PresencePenalty,
	}

	data, err := json.Marshal(req)
	if err != nil {
		return nil, err
	}

	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, b.baseURL+"/v1/chat/completions", bytes.NewReader(data))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if b.apiKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+b.apiKey)
	}
	for k, v := range b.extraHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := b.client.Do(httpReq)
	if err != nil {
		return nil, err
	}
	if resp.StatusCode == http.StatusTooManyRequests {
		body, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		retryAfter := parseRetryAfter(resp.Header.Get("Retry-After"))
		return nil, &RateLimitError{RetryAfter: retryAfter, Message: string(body)}
	}
	if resp.StatusCode != http.StatusOK {
		body, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("openai HTTP %d: %s", resp.StatusCode, body)
	}

	ch := make(chan StreamEvent, 8)

	go func() {
		defer close(ch)
		defer resp.Body.Close()

		scanner := bufio.NewScanner(resp.Body)
		scanner.Buffer(make([]byte, 0, 64*1024), 4*1024*1024)

		// Accumulate tool call fragments by index (OpenAI sends args incrementally).
		type tcAccum struct {
			id, name, args string
		}
		accum := make(map[int]*tcAccum)

		// finishReason and usage arrive in separate trailing chunks.
		var finishReason string
		var usage Usage

		for scanner.Scan() {
			line := strings.TrimSpace(scanner.Text())
			if line == "" {
				continue
			}
			if line == "data: [DONE]" {
				ev := StreamEvent{
					Done:       true,
					StopReason: finishReason,
					Usage:      usage,
				}
				for _, a := range accum {
					ev.ToolCalls = append(ev.ToolCalls, ToolCall{
						ID:        a.id,
						Name:      a.name,
						Arguments: json.RawMessage(a.args),
					})
				}
				ch <- ev
				return
			}
			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			var wire openAIStreamResponse
			if err := json.Unmarshal([]byte(line[6:]), &wire); err != nil {
				ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream decode: %v", err)}
				return
			}

			if wire.Usage.PromptTokens > 0 {
				usage.InputTokens = wire.Usage.PromptTokens
			}
			if wire.Usage.CompletionTokens > 0 {
				usage.OutputTokens = wire.Usage.CompletionTokens
			}

			if len(wire.Choices) == 0 {
				continue
			}

			choice := wire.Choices[0]
			if choice.FinishReason != "" {
				finishReason = choice.FinishReason
			}
			if choice.Delta.Content != "" {
				ch <- StreamEvent{Content: choice.Delta.Content}
			}

			for _, dtc := range choice.Delta.ToolCalls {
				idx := dtc.Index
				if _, ok := accum[idx]; !ok {
					accum[idx] = &tcAccum{id: dtc.ID, name: dtc.Function.Name}
				}
				a := accum[idx]
				if dtc.ID != "" {
					a.id = dtc.ID
				}
				if dtc.Function.Name != "" {
					a.name = dtc.Function.Name
				}
				a.args += dtc.Function.Arguments
			}
		}
		if err := scanner.Err(); err != nil {
			ch <- StreamEvent{Done: true, StopReason: fmt.Sprintf("stream read: %v", err)}
		}
	}()

	return ch, nil
}

// parseRetryAfter parses the Retry-After header value, which may be an integer
// number of seconds or an HTTP-date. Returns zero if the header is absent or
// unparseable.
func parseRetryAfter(header string) time.Duration {
	if header == "" {
		return 0
	}
	header = strings.TrimSpace(header)
	if secs, err := strconv.Atoi(header); err == nil {
		return time.Duration(secs) * time.Second
	}
	if t, err := http.ParseTime(header); err == nil {
		if d := time.Until(t); d > 0 {
			return d
		}
	}
	return 0
}