package main

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net/http"
	"strings"
	"time"
)

type ModelConfig struct {
	ID      string `yaml:"id"`
	OwnedBy string `yaml:"owned_by"`
}

// Config holds the application configuration
type Config struct {
	Port        int          `yaml:"port"`
	UpstreamURL string       `yaml:"upstream_url"`
	Models      []ModelConfig `yaml:"models"`
	Temperature *float64     `yaml:"temperature,omitempty"`
	// Retry configuration
	MaxRetries      int `yaml:"max_retries,omitempty"` // Maximum retry attempts for 429/503
	RetryBaseDelayMs int `yaml:"retry_base_delay_ms,omitempty"` // Base delay in ms (exponential backoff)
}

var config *Config

// httpClient is a shared HTTP client for all upstream requests
var httpClient = &http.Client{Timeout: 300 * time.Second}

// maxBodySize limits request body size to prevent memory exhaustion
const maxBodySize = 10 << 20 // 10MB

// blockedHeaders are headers that should never be forwarded to upstream
// for security/privacy reasons. These headers could leak internal URLs,
// session information, or other sensitive data.
var blockedHeaders = map[string]bool{
	"Referer":         true, // Don't leak internal URLs to external API
	"Cookie":          true, // Don't forward session cookies
	"Authorization":   true, // Already extracted and sent as x-api-key
	"X-Forwarded-For": true, // Don't leak client IP
	"X-Real-Ip":       true, // Don't leak client IP
	"X-Forwarded-Host": true, // Don't leak internal hostnames
}

// ClaudeCodeHeaders returns the headers to mimic claude-code CLI
func ClaudeCodeHeaders(apiKey, sessionID string) map[string]string {
	return map[string]string{
		"User-Agent":               "claude-cli/1.0.18 (pro, cli)",
		"x-api-key":                apiKey,
		"x-app":                    "cli",
		"anthropic-version":        "2023-06-01",
		"anthropic-beta":           "claude-code-20250219,interleaved-thinking-2025-05-14,prompt-caching-scope-2026-01-05,context-management-2025-06-27,redact-thinking-2026-02-12",
		"X-Claude-Code-Session-Id": sessionID,
		"content-type":             "application/json",
	}
}

func handleModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodGet {
		writeError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", "method_not_allowed")
		return
	}

	models := make([]map[string]interface{}, len(config.Models))
	for i, model := range config.Models {
		models[i] = map[string]interface{}{
			"id":       model.ID,
			"object":   "model",
			"created":  1234567890,
			"owned_by": model.OwnedBy,
		}
	}

	response := map[string]interface{}{
		"object": "list",
		"data":   models,
	}

	w.Header().Set("content-type", "application/json")
	json.NewEncoder(w).Encode(response)
}

func handleChatCompletions(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		writeError(w, http.StatusMethodNotAllowed, "Method not allowed", "invalid_request_error", "method_not_allowed")
		return
	}

	// Extract Bearer token
	authHeader := r.Header.Get("Authorization")
	if authHeader == "" || !strings.HasPrefix(authHeader, "Bearer ") {
		writeError(w, http.StatusUnauthorized, "Missing or invalid Authorization header", "authentication_error", "missing_authorization")
		return
	}
	apiKey := strings.TrimPrefix(authHeader, "Bearer ")

	// Read body (with size limit to prevent memory exhaustion)
	r.Body = http.MaxBytesReader(w, r.Body, maxBodySize)
	body, err := io.ReadAll(r.Body)
	if err != nil {
		// Check if this is a max bytes error
		if err.Error() == "http: request body too large" {
			writeError(w, http.StatusRequestEntityTooLarge, "Request body exceeds maximum size limit (10MB)", "invalid_request_error", "body_too_large")
			return
		}
		writeError(w, http.StatusBadRequest, "Failed to read request body", "invalid_request_error", "body_read_error")
		return
	}

	// Decode request
	var req ChatCompletionRequest
	if err := json.Unmarshal(body, &req); err != nil {
		writeError(w, http.StatusBadRequest, "Invalid JSON in request body", "invalid_request_error", "json_decode_error")
		return
	}

	// Get session ID from context (set by main)
	sessionID, ok := r.Context().Value(sessionIDKey).(string)
	if !ok {
		writeError(w, http.StatusInternalServerError, "Internal server error", "internal_error", "session_missing")
		return
	}

	// Convert to Anthropic format — always non-streaming to upstream
	// (ZAI's streaming returns empty for GLM models)
	anthropicReq := ConvertOpenAIRequest(&req, config.Temperature)
	anthropicReq.Stream = false

	reqBody, _ := json.Marshal(anthropicReq)
	log.Printf("[debug] Sending to upstream %s, model=%s, body_size=%d bytes", config.UpstreamURL, req.Model, len(reqBody))

	// Non-streaming request to upstream
	upstreamResp, err := callUpstream(anthropicReq, apiKey, sessionID)
	if err != nil {
		writeError(w, http.StatusBadGateway, fmt.Sprintf("Upstream request failed: %v", err), "upstream_error", "proxy_error")
		return
	}
	defer upstreamResp.Body.Close()

	if upstreamResp.StatusCode != http.StatusOK {
		respBody, _ := io.ReadAll(upstreamResp.Body)
		log.Printf("[debug] Upstream error status %d, body_size=%d bytes", upstreamResp.StatusCode, len(respBody))
		writeError(w, http.StatusBadGateway, fmt.Sprintf("Upstream returned error: %s", string(respBody)), "upstream_error", fmt.Sprintf("status_%d", upstreamResp.StatusCode))
		return
	}

	// Read the full Anthropic response
	respBody, err := io.ReadAll(upstreamResp.Body)
	if err != nil {
		writeError(w, http.StatusBadGateway, "Failed to read upstream response", "upstream_error", "body_read_error")
		return
	}
	log.Printf("[debug] Upstream response received, body_size=%d bytes", len(respBody))

	var anthropicResp AnthropicResponse
	if err := json.Unmarshal(respBody, &anthropicResp); err != nil {
		writeError(w, http.StatusBadGateway, "Failed to parse upstream response", "upstream_error", "json_decode_error")
		return
	}

	isStream := req.Stream != nil && *req.Stream

	if isStream {
		// Convert the non-streaming response to SSE chunks for the client
		w.Header().Set("content-type", "text/event-stream")
		w.Header().Set("cache-control", "no-cache")
		w.Header().Set("connection", "keep-alive")

		created := time.Now().Unix()
		chunkID := "chatcmpl-" + randomString(8)

		// Extract text content
		var textContent string
		for _, block := range anthropicResp.Content {
			if block.Type == "text" {
				textContent += block.Text
			}
		}

		// Send text as chunks (simulate streaming)
		if textContent != "" {
			chunk := StreamChunk{
				ID:      chunkID,
				Object:  "chat.completion.chunk",
				Created: created,
				Model:   req.Model,
				Choices: []StreamChoice{
					{
						Index: 0,
						Delta: Delta{
							Role:    "assistant",
							Content: textContent,
						},
					},
				},
			}
			data, _ := json.Marshal(chunk)
			fmt.Fprintf(w, "data: %s\n\n", data)
			w.(http.Flusher).Flush()
		}

		// Send finish chunk
		finishChunk := StreamChunk{
			ID:      chunkID,
			Object:  "chat.completion.chunk",
			Created: created,
			Model:   req.Model,
			Choices: []StreamChoice{
				{
					Index:        0,
					Delta:        Delta{},
					FinishReason: mapStopReason(anthropicResp.StopReason),
				},
			},
		}
		data, _ := json.Marshal(finishChunk)
		fmt.Fprintf(w, "data: %s\n\n", data)
		w.(http.Flusher).Flush()

		fmt.Fprintf(w, "data: [DONE]\n\n")
		w.(http.Flusher).Flush()
	} else {
		// Non-streaming: convert directly
		openAIResp := ConvertAnthropicResponse(&anthropicResp, req.Model)
		w.Header().Set("content-type", "application/json")
		json.NewEncoder(w).Encode(openAIResp)
	}
}

// retryableStatuses are HTTP status codes that should trigger a retry
var retryableStatuses = map[int]bool{
	429: true, // Rate limit
	503: true, // Service unavailable
}

func callUpstream(req *AnthropicRequest, apiKey, sessionID string) (*http.Response, error) {
	bodyBytes, err := json.Marshal(req)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal request: %w", err)
	}

	upstreamURL := config.UpstreamURL + "/v1/messages"

	// Get retry config with defaults
	maxRetries := config.MaxRetries
	if maxRetries == 0 {
		maxRetries = 3
	}
	baseDelayMs := config.RetryBaseDelayMs
	if baseDelayMs == 0 {
		baseDelayMs = 1000
	}

	var lastResp *http.Response
	var lastErr error

	for attempt := 0; attempt <= maxRetries; attempt++ {
		if attempt > 0 {
			delay := time.Duration(baseDelayMs*(1<<(attempt-1))) * time.Millisecond
			// Add jitter (±50%) to avoid thundering herd
			jitter := time.Duration(rand.Intn(int(delay))) / 2
			log.Printf("[retry] Attempt %d/%d after %v (jitter: %v)", attempt, maxRetries, delay, jitter)
			time.Sleep(delay + jitter)
		}

		httpReq, err := http.NewRequest(http.MethodPost, upstreamURL, strings.NewReader(string(bodyBytes)))
		if err != nil {
			lastErr = fmt.Errorf("failed to create request: %w", err)
			continue
		}

		headers := ClaudeCodeHeaders(apiKey, sessionID)
		for k, v := range headers {
			httpReq.Header.Set(k, v)
		}

		resp, err := httpClient.Do(httpReq)
		if err != nil {
			lastErr = err
			log.Printf("[retry] Request failed: %v", err)
			continue
		}

		// Check if status is retryable
		if !retryableStatuses[resp.StatusCode] {
			log.Printf("[retry] Success or non-retryable status %d", resp.StatusCode)
			return resp, nil
		}

		// Close previous response body if exists
		if lastResp != nil {
			lastResp.Body.Close()
		}

		// Read and log the error body for debugging
		respBody, _ := io.ReadAll(resp.Body)
		resp.Body.Close()
		log.Printf("[retry] Retryable status %d, body: %s", resp.StatusCode, string(respBody))

		lastResp = resp
		lastErr = fmt.Errorf("upstream returned status %d", resp.StatusCode)
	}

	// All retries exhausted
	return nil, lastErr
}

func writeError(w http.ResponseWriter, code int, message, errType, errCode string) {
	w.Header().Set("content-type", "application/json")
	w.WriteHeader(code)
	resp := map[string]interface{}{
		"error": map[string]string{
			"message": message,
			"type":    errType,
			"code":    errCode,
		},
	}
	json.NewEncoder(w).Encode(resp)
}

func randomString(n int) string {
	const letters = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789"
	b := make([]byte, n)
	for i := range b {
		b[i] = letters[rand.Intn(len(letters))]
	}
	return string(b)
}
