package providers

import (
	"switchboard-core/capabilities"
	"bufio"
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// OpenAIProvider handles any OpenAI-compatible API.
type OpenAIProvider struct{}

func (p *OpenAIProvider) ID() string { return "openai" }

// ── Chat Completion (non-streaming) ─────────

func (p *OpenAIProvider) ChatCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (*CompletionResponse, error) {
	req.Stream = false

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}
	defer body.Close()

	var resp openaiChatResponse
	if err := json.NewDecoder(body).Decode(&resp); err != nil {
		return nil, fmt.Errorf("decode response: %w", err)
	}

	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("no choices in response")
	}

	// Extract content string from response (responses are always string, not array)
	var contentStr string
	if s, ok := resp.Choices[0].Message.Content.(string); ok {
		contentStr = s
	}

	result := &CompletionResponse{
		Content:      contentStr,
		Model:        resp.Model,
		FinishReason: resp.Choices[0].FinishReason,
		InputTokens:  resp.Usage.PromptTokens,
		OutputTokens: resp.Usage.CompletionTokens,
	}

	// Cache tokens (OpenAI, Venice report these in prompt_tokens_details)
	if resp.Usage.PromptTokensDetails != nil {
		result.CacheReadTokens = resp.Usage.PromptTokensDetails.CachedTokens
	}

	// Extract tool calls if present
	if len(resp.Choices[0].Message.ToolCalls) > 0 {
		for _, tc := range resp.Choices[0].Message.ToolCalls {
			result.ToolCalls = append(result.ToolCalls, ToolCall{
				ID:   tc.ID,
				Type: tc.Type,
				Function: FunctionCall{
					Name:      tc.Function.Name,
					Arguments: tc.Function.Arguments,
				},
			})
		}
		if result.FinishReason == "" {
			result.FinishReason = "tool_calls"
		}
	}

	return result, nil
}

// ── Stream Completion ───────────────────────

func (p *OpenAIProvider) StreamCompletion(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (<-chan StreamEvent, error) {
	req.Stream = true

	body, err := p.doRequest(ctx, cfg, req)
	if err != nil {
		return nil, err
	}

	ch := make(chan StreamEvent, 64)

	go func() {
		defer close(ch)
		defer body.Close()

		// Accumulate tool calls across stream chunks
		toolCallMap := map[int]*ToolCall{} // index → accumulated call

		// Usage arrives in a separate chunk (or on the final chunk).
		// OpenAI protocol order: content → finish_reason → usage → [DONE]
		// The usage chunk has choices=[] so we must hold the finish event
		// until usage arrives, otherwise tokens are always 0.
		var streamUsage *openaiUsage
		var pendingFinish *StreamEvent // deferred until usage arrives

		scanner := bufio.NewScanner(body)
		for scanner.Scan() {
			line := scanner.Text()

			if !strings.HasPrefix(line, "data: ") {
				continue
			}

			data := strings.TrimPrefix(line, "data: ")
			if data == "[DONE]" {
				// Flush any pending finish event (usage may or may not have arrived)
				if pendingFinish != nil {
					ch <- *pendingFinish
				} else {
					ch <- StreamEvent{Done: true}
				}
				return
			}

			var chunk openaiStreamChunk
			if err := json.Unmarshal([]byte(data), &chunk); err != nil {
				continue
			}

			// Capture usage — may arrive in a separate chunk with empty choices
			if chunk.Usage != nil {
				streamUsage = chunk.Usage
				// If finish already arrived, attach usage and flush immediately
				if pendingFinish != nil {
					pendingFinish.InputTokens = streamUsage.PromptTokens
					pendingFinish.OutputTokens = streamUsage.CompletionTokens
					if streamUsage.PromptTokensDetails != nil {
						pendingFinish.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
					}
					ch <- *pendingFinish
					return // Done — [DONE] will be handled by defer body.Close()
				}
			}

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

			choice := chunk.Choices[0]

			// Accumulate tool call deltas
			for _, tc := range choice.Delta.ToolCalls {
				idx := 0
				if tc.Index != nil {
					idx = *tc.Index
				}

				if existing, ok := toolCallMap[idx]; ok {
					// Append arguments fragment
					existing.Function.Arguments += tc.Function.Arguments
				} else {
					// First chunk for this tool — has ID and name
					toolCallMap[idx] = &ToolCall{
						ID:   tc.ID,
						Type: "function",
						Function: FunctionCall{
							Name:      tc.Function.Name,
							Arguments: tc.Function.Arguments,
						},
					}
				}
			}

			ev := StreamEvent{
				Delta:        choice.Delta.Content,
				Reasoning:    choice.Delta.ReasoningContent,
				Model:        chunk.Model,
				FinishReason: choice.FinishReason,
			}

			if ev.FinishReason != "" {
				ev.Done = true
				// Attach accumulated tool calls on final event
				if ev.FinishReason == "tool_calls" && len(toolCallMap) > 0 {
					for _, tc := range toolCallMap {
						ev.ToolCalls = append(ev.ToolCalls, *tc)
					}
					// Tool calls need to flush immediately — the tool loop
					// needs the event to start executing tools
					if streamUsage != nil {
						ev.InputTokens = streamUsage.PromptTokens
						ev.OutputTokens = streamUsage.CompletionTokens
						if streamUsage.PromptTokensDetails != nil {
							ev.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
						}
					}
					ch <- ev
					continue
				}
				// Normal finish — defer until usage chunk arrives
				if streamUsage != nil {
					// Usage already captured (rare: same chunk or earlier)
					ev.InputTokens = streamUsage.PromptTokens
					ev.OutputTokens = streamUsage.CompletionTokens
					if streamUsage.PromptTokensDetails != nil {
						ev.CacheReadTokens = streamUsage.PromptTokensDetails.CachedTokens
					}
					ch <- ev
					continue
				}
				// Hold — usage chunk hasn't arrived yet
				pendingFinish = &ev
				continue
			}

			ch <- ev
		}

		if err := scanner.Err(); err != nil {
			ch <- StreamEvent{Error: err}
		} else if pendingFinish != nil {
			// Scanner ended without [DONE] — flush pending finish
			ch <- *pendingFinish
		}
	}()

	return ch, nil
}

// ── List Models ─────────────────────────────

func (p *OpenAIProvider) ListModels(ctx context.Context, cfg ProviderConfig) ([]Model, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/models"

	httpReq, err := http.NewRequestWithContext(ctx, "GET", endpoint, nil)
	if err != nil {
		return nil, err
	}
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := cfg.Client().Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("list models: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("list models: HTTP %d: %s", resp.StatusCode, string(b))
	}

	var result openaiModelsResponse
	if err := json.NewDecoder(resp.Body).Decode(&result); err != nil {
		return nil, fmt.Errorf("decode models: %w", err)
	}

	out := make([]Model, 0, len(result.Data))
	for _, m := range result.Data {
		// Heuristic inference from model ID
		caps := capabilities.InferCapabilities(m.ID)
		// Use context_length from API if available and we don't have it
		if m.ContextLength > 0 && caps.MaxContext == 0 {
			caps.MaxContext = m.ContextLength
		}

		out = append(out, Model{
			ID:           m.ID,
			Type:         m.Type, // passthrough from API if present (empty → "chat" at sync time)
			OwnedBy:      m.OwnedBy,
			Capabilities: caps,
		})
	}
	return out, nil
}

// ── Embeddings ─────────────────────────────

func (p *OpenAIProvider) Embed(ctx context.Context, cfg ProviderConfig, req EmbeddingRequest) (*EmbeddingResponse, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/embeddings"

	oaiReq := openaiEmbeddingRequest{
		Model: req.Model,
		Input: req.Input,
	}
	body, err := json.Marshal(oaiReq)
	if err != nil {
		return nil, fmt.Errorf("marshal embedding request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := cfg.Client().Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("embedding request: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		return nil, fmt.Errorf("embedding error: HTTP %d: %s", resp.StatusCode, string(b))
	}

	var oaiResp openaiEmbeddingResponse
	if err := json.NewDecoder(resp.Body).Decode(&oaiResp); err != nil {
		return nil, fmt.Errorf("decode embedding response: %w", err)
	}

	result := &EmbeddingResponse{
		Model:       oaiResp.Model,
		InputTokens: oaiResp.Usage.PromptTokens,
		Embeddings:  make([][]float64, len(oaiResp.Data)),
	}
	for i, d := range oaiResp.Data {
		result.Embeddings[i] = d.Embedding
	}
	return result, nil
}

// ── HTTP Layer ──────────────────────────────

func (p *OpenAIProvider) doRequest(ctx context.Context, cfg ProviderConfig, req CompletionRequest) (io.ReadCloser, error) {
	endpoint := strings.TrimSuffix(cfg.Endpoint, "/") + "/chat/completions"

	// Build OpenAI-format request
	oaiReq := openaiChatRequest{
		Model:    req.Model,
		Stream:   req.Stream,
		Messages: make([]openaiMessage, 0, len(req.Messages)),
	}
	if req.Stream {
		oaiReq.StreamOptions = &openaiStreamOptions{IncludeUsage: true}
	}
	if req.MaxTokens > 0 {
		oaiReq.MaxTokens = req.MaxTokens
	}
	if req.Temperature != nil {
		oaiReq.Temperature = req.Temperature
	}
	if req.TopP != nil {
		oaiReq.TopP = req.TopP
	}

	// Convert tools
	for _, t := range req.Tools {
		oaiReq.Tools = append(oaiReq.Tools, openaiToolDef{
			Type: t.Type,
			Function: openaiToolDefFunction{
				Name:        t.Function.Name,
				Description: t.Function.Description,
				Parameters:  t.Function.Parameters,
			},
		})
	}

	// Convert messages — handle all roles including tool and multimodal
	for _, m := range req.Messages {
		oaiMsg := openaiMessage{
			Role: m.Role,
		}

		// Build content: multimodal parts or plain string
		if len(m.ContentParts) > 0 && m.Role == "user" {
			// Multimodal: convert to OpenAI content array format
			parts := make([]openaiContentPart, 0, len(m.ContentParts))
			for _, p := range m.ContentParts {
				switch p.Type {
				case "image_url":
					if p.ImageURL != nil {
						detail := p.ImageURL.Detail
						if detail == "" {
							detail = "auto"
						}
						parts = append(parts, openaiContentPart{
							Type: "image_url",
							ImageURL: &openaiImageURL{
								URL:    p.ImageURL.URL,
								Detail: detail,
							},
						})
					}
				default: // "text", "document" → text
					parts = append(parts, openaiContentPart{
						Type: "text",
						Text: p.Text,
					})
				}
			}
			oaiMsg.Content = parts
		} else {
			oaiMsg.Content = m.Content
		}

		// Assistant messages with tool calls
		if len(m.ToolCalls) > 0 {
			for _, tc := range m.ToolCalls {
				oaiMsg.ToolCalls = append(oaiMsg.ToolCalls, openaiToolCall{
					ID:   tc.ID,
					Type: tc.Type,
					Function: openaiToolCallFunction{
						Name:      tc.Function.Name,
						Arguments: tc.Function.Arguments,
					},
				})
			}
		}

		// Tool result messages
		if m.Role == "tool" {
			oaiMsg.ToolCallID = m.ToolCallID
			oaiMsg.Name = m.Name
		}

		// Named messages: pass through for speaker attribution (v0.23.0)
		// OpenAI supports "name" on any role to distinguish participants
		if m.Name != "" && m.Role != "tool" {
			oaiMsg.Name = m.Name
		}

		oaiReq.Messages = append(oaiReq.Messages, oaiMsg)
	}

	body, err := json.Marshal(oaiReq)
	if err != nil {
		return nil, fmt.Errorf("marshal request: %w", err)
	}

	// Merge hook-supplied extra fields into wire JSON (v0.22.1)
	if len(req.ExtraBody) > 0 {
		body, err = mergeExtraBody(body, req.ExtraBody)
		if err != nil {
			return nil, fmt.Errorf("merge extra body: %w", err)
		}
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST", endpoint, bytes.NewReader(body))
	if err != nil {
		return nil, err
	}
	httpReq.Header.Set("Content-Type", "application/json")
	if cfg.APIKey != "" {
		httpReq.Header.Set("Authorization", "Bearer "+cfg.APIKey)
	}
	// Inject provider-level custom headers (OpenRouter, Venice, etc.)
	for k, v := range cfg.CustomHeaders {
		httpReq.Header.Set(k, v)
	}

	resp, err := cfg.Client().Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("provider request: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		b, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		return nil, fmt.Errorf("provider error: HTTP %d: %s", resp.StatusCode, string(b))
	}

	return resp.Body, nil
}

// ── OpenAI Wire Types ───────────────────────

type openaiMessage struct {
	Role         string              `json:"role"`
	Content      interface{}         `json:"content,omitempty"`      // string or []openaiContentPart
	ToolCalls    []openaiToolCall    `json:"tool_calls,omitempty"`   // assistant requesting tools
	ToolCallID   string              `json:"tool_call_id,omitempty"` // role=tool result
	Name         string              `json:"name,omitempty"`         // speaker name (tool or persona attribution)
}

// openaiContentPart represents a single element in a multimodal content array.
// OpenAI format: [{"type":"text","text":"..."}, {"type":"image_url","image_url":{"url":"..."}}]
type openaiContentPart struct {
	Type     string          `json:"type"`
	Text     string          `json:"text,omitempty"`
	ImageURL *openaiImageURL `json:"image_url,omitempty"`
}

type openaiImageURL struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"`
}

type openaiToolCallFunction struct {
	Name      string `json:"name,omitempty"`
	Arguments string `json:"arguments,omitempty"`
}

type openaiToolCall struct {
	Index    *int                   `json:"index,omitempty"` // present in streaming deltas
	ID       string                 `json:"id,omitempty"`
	Type     string                 `json:"type,omitempty"` // "function"
	Function openaiToolCallFunction `json:"function"`
}

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

type openaiToolDef struct {
	Type     string               `json:"type"` // "function"
	Function openaiToolDefFunction `json:"function"`
}

type openaiChatRequest struct {
	Model         string               `json:"model"`
	Messages      []openaiMessage      `json:"messages"`
	MaxTokens     int                  `json:"max_tokens,omitempty"`
	Temperature   *float64             `json:"temperature,omitempty"`
	TopP          *float64             `json:"top_p,omitempty"`
	Stream        bool                 `json:"stream"`
	StreamOptions *openaiStreamOptions `json:"stream_options,omitempty"`
	Tools         []openaiToolDef      `json:"tools,omitempty"`
}

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

type openaiChatResponse struct {
	Model   string `json:"model"`
	Choices []struct {
		Message      openaiMessage `json:"message"`
		FinishReason string        `json:"finish_reason"`
	} `json:"choices"`
	Usage openaiUsage `json:"usage"`
}

type openaiUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	PromptTokensDetails *struct {
		CachedTokens int `json:"cached_tokens"`
	} `json:"prompt_tokens_details,omitempty"`
	CompletionTokensDetails *struct {
		ReasoningTokens int `json:"reasoning_tokens"`
	} `json:"completion_tokens_details,omitempty"`
}

type openaiStreamChunk struct {
	Model   string `json:"model"`
	Choices []struct {
		Delta struct {
			Content          string           `json:"content"`
			ReasoningContent string           `json:"reasoning_content"`
			ToolCalls        []openaiToolCall `json:"tool_calls,omitempty"`
		} `json:"delta"`
		FinishReason string `json:"finish_reason"`
	} `json:"choices"`
	Usage *openaiUsage `json:"usage,omitempty"`
}

type openaiModelsResponse struct {
	Data []struct {
		ID            string `json:"id"`
		OwnedBy       string `json:"owned_by"`
		Type          string `json:"type,omitempty"`              // Some APIs (e.g. Venice-compat) include model type
		ContextLength int    `json:"context_length,omitempty"` // Ollama, OpenRouter
	} `json:"data"`
}

// ── Embedding Wire Types ──────────────────

type openaiEmbeddingRequest struct {
	Model string   `json:"model"`
	Input []string `json:"input"`
}

type openaiEmbeddingResponse struct {
	Data []struct {
		Embedding []float64 `json:"embedding"`
		Index     int       `json:"index"`
	} `json:"data"`
	Model string `json:"model"`
	Usage struct {
		PromptTokens int `json:"prompt_tokens"`
	} `json:"usage"`
}