package backend

import (
	"bytes"
	"context"
	"encoding/base64"
	"encoding/json"
	"errors"
	"fmt"
	"image"
	"image/color"
	"image/png"
	"io"
	"net"
	"net/http"
	"strings"
	"time"
)

const (
	defaultImageExtractPrompt = `请分析这张图片，输出用于 Drive9 文件搜索的结构化 JSON。只输出一个 JSON 对象，不要 Markdown，不要解释。基础字段为：
{
  "caption_zh": "中文一句话摘要，<=120字",
  "description_zh": "中文详细描述，覆盖主体、场景、动作、颜色、构图、风格、时间/季节/地点线索，<=1200字",
  "caption_en": "one concise English caption, <=30 words",
  "description_en": "English description covering subjects, scene, action, colors, composition, style, time/season/location clues, <=260 words",
  "ocr_text": ["图中可见文字；没有则空数组；最多100条，每条<=160字"],
  "tags_zh": ["中文搜索标签，5-30个，短词，不要泛词"],
  "tags_en": ["English search tags, 5-30, lowercase short phrases, no generic words like image/photo/picture/ocr/text"],
  "search_queries_zh": ["中文自然搜索短语，3-12个"],
  "search_queries_en": ["English natural search phrases, 3-12"]
}
要求：
- 中英文都要覆盖可搜索信息，但不要逐字翻译造成重复堆砌。
- 标签优先使用具体物体、场景、风格、颜色、季节、地点、动作、可见文字主题。
- 若不确定不要编造；没有 OCR 就使用 []。
- 整个 JSON 控制在 16000 字符以内。`

	// DefaultOpenAIImageExtractMaxTokens leaves enough completion budget for
	// structured bilingual extraction while the writeback path still enforces
	// Drive9's embedding input safety limits.
	DefaultOpenAIImageExtractMaxTokens = 4096

	imageProviderValidationMaxTokens = 16
	imageProviderValidationPrompt    = "Reply with one short word describing this image."
)

// ImageProviderFailureCategory is a low-cardinality provider failure suitable
// for logs, metrics, API responses, and fixed semantic task error codes.
type ImageProviderFailureCategory string

const (
	ImageProviderFailureAuthRejected    ImageProviderFailureCategory = "auth_rejected"
	ImageProviderFailureNotFound        ImageProviderFailureCategory = "not_found"
	ImageProviderFailureRequestRejected ImageProviderFailureCategory = "request_rejected"
	ImageProviderFailureRateLimited     ImageProviderFailureCategory = "rate_limited"
	ImageProviderFailureTimeout         ImageProviderFailureCategory = "timeout"
	ImageProviderFailureNetworkError    ImageProviderFailureCategory = "network_error"
	ImageProviderFailureProviderError   ImageProviderFailureCategory = "provider_error"
	ImageProviderFailureInvalidResponse ImageProviderFailureCategory = "invalid_response"
)

type imageProviderError struct {
	category ImageProviderFailureCategory
	status   int
	cause    error
}

func (e *imageProviderError) Error() string {
	reason := map[ImageProviderFailureCategory]string{
		ImageProviderFailureAuthRejected:    "authentication rejected",
		ImageProviderFailureNotFound:        "endpoint or model not found",
		ImageProviderFailureRequestRejected: "request rejected; verify that the model supports image input",
		ImageProviderFailureRateLimited:     "provider rate limited the request",
		ImageProviderFailureTimeout:         "provider request timed out",
		ImageProviderFailureNetworkError:    "provider network request failed",
		ImageProviderFailureProviderError:   "provider is unavailable",
		ImageProviderFailureInvalidResponse: "provider returned an invalid response",
	}[e.category]
	if reason == "" {
		reason = "provider request failed"
	}
	if e.status > 0 {
		return fmt.Sprintf("%s (HTTP %d)", reason, e.status)
	}
	return reason
}

func (e *imageProviderError) Unwrap() error { return e.cause }

// ImageProviderFailure returns the stable category and HTTP status carried by
// a provider failure. It never returns provider-supplied response text.
func ImageProviderFailure(err error) (ImageProviderFailureCategory, int, bool) {
	var providerErr *imageProviderError
	if errors.As(err, &providerErr) {
		return providerErr.category, providerErr.status, true
	}
	var audioErr *AudioExtractAPIError
	if errors.As(err, &audioErr) {
		classified := imageProviderHTTPError(audioErr.StatusCode)
		if errors.As(classified, &providerErr) {
			return providerErr.category, providerErr.status, true
		}
	}
	return "", 0, false
}

func imageProviderHTTPError(status int) error {
	category := ImageProviderFailureRequestRejected
	switch {
	case status == http.StatusUnauthorized || status == http.StatusForbidden:
		category = ImageProviderFailureAuthRejected
	case status == http.StatusNotFound:
		category = ImageProviderFailureNotFound
	case status == http.StatusRequestTimeout || status == http.StatusGatewayTimeout:
		category = ImageProviderFailureTimeout
	case status == http.StatusTooManyRequests:
		category = ImageProviderFailureRateLimited
	case status >= http.StatusInternalServerError:
		category = ImageProviderFailureProviderError
	}
	return &imageProviderError{category: category, status: status}
}

func imageProviderTransportError(err error) error {
	category := ImageProviderFailureNetworkError
	var netErr net.Error
	if errors.Is(err, errTenantExtractorEndpointPolicy) {
		category = ImageProviderFailureRequestRejected
	} else if errors.Is(err, context.DeadlineExceeded) || errors.As(err, &netErr) && netErr.Timeout() {
		category = ImageProviderFailureTimeout
	}
	return &imageProviderError{category: category, cause: err}
}

func invalidImageProviderResponse(status int, cause error) error {
	return &imageProviderError{category: ImageProviderFailureInvalidResponse, status: status, cause: cause}
}

// OpenAIImageTextExtractorConfig configures an OpenAI-compatible vision endpoint.
// This works with providers that expose the /v1/chat/completions API surface.
type OpenAIImageTextExtractorConfig struct {
	BaseURL string
	APIKey  string
	Model   string
	// Prompt is the fully composed extraction prompt. Callers should build it
	// with ComposeImageExtractPrompt so the attributes suffix is applied.
	Prompt    string
	MaxTokens int
	Timeout   time.Duration
	// Client must enforce tenant endpoint policy (for example, by using
	// NewTenantExtractorHTTPClient) whenever BaseURL is tenant-controlled. A
	// nil client is reserved for trusted operator/process configuration, which
	// may intentionally use private endpoints or environment proxies.
	Client *http.Client
}

type OpenAIImageTextExtractor struct {
	endpoint  string
	apiKey    string
	model     string
	prompt    string
	maxTokens int
	client    *http.Client
}

func NewOpenAIImageTextExtractor(cfg OpenAIImageTextExtractorConfig) (*OpenAIImageTextExtractor, error) {
	base := strings.TrimRight(cfg.BaseURL, "/")
	if base == "" {
		return nil, fmt.Errorf("openai extractor base url is required")
	}
	if cfg.APIKey == "" {
		return nil, fmt.Errorf("openai extractor api key is required")
	}
	if cfg.Model == "" {
		return nil, fmt.Errorf("openai extractor model is required")
	}
	var endpoint string
	if strings.HasSuffix(base, "/v1") {
		endpoint = base + "/chat/completions"
	} else {
		endpoint = base + "/v1/chat/completions"
	}
	if cfg.MaxTokens <= 0 {
		cfg.MaxTokens = DefaultOpenAIImageExtractMaxTokens
	}
	client := cfg.Client
	if client == nil {
		timeout := cfg.Timeout
		if timeout <= 0 {
			timeout = 30 * time.Second
		}
		client = &http.Client{Timeout: timeout}
	}
	return &OpenAIImageTextExtractor{
		endpoint:  endpoint,
		apiKey:    cfg.APIKey,
		model:     cfg.Model,
		prompt:    cfg.Prompt,
		maxTokens: cfg.MaxTokens,
		client:    client,
	}, nil
}

// ValidateOpenAIImageProvider sends one small real vision request using the
// same extractor and SSRF-safe HTTP client used by tenant runtime extraction.
func ValidateOpenAIImageProvider(ctx context.Context, cfg OpenAIImageTextExtractorConfig) error {
	if cfg.Client == nil {
		cfg.Client = NewTenantExtractorHTTPClient(cfg.Timeout)
	}
	cfg.Prompt = imageProviderValidationPrompt
	cfg.MaxTokens = imageProviderValidationMaxTokens
	extractor, err := NewOpenAIImageTextExtractor(cfg)
	if err != nil {
		return err
	}
	return validateOpenAIImageProvider(ctx, extractor)
}

func validateOpenAIImageProvider(ctx context.Context, extractor ImageTextExtractor) error {
	img := image.NewRGBA(image.Rect(0, 0, 32, 32))
	for y := 0; y < 32; y++ {
		for x := 0; x < 32; x++ {
			img.SetRGBA(x, y, color.RGBA{R: 255, G: 255, B: 255, A: 255})
		}
	}
	var data bytes.Buffer
	if err := png.Encode(&data, img); err != nil {
		return fmt.Errorf("encode provider validation image: %w", err)
	}
	_, _, err := extractor.ExtractImageText(ctx, ImageExtractRequest{
		ContentType: "image/png",
		Data:        data.Bytes(),
	})
	return err
}

func (e *OpenAIImageTextExtractor) ExtractImageText(ctx context.Context, req ImageExtractRequest) (string, ImageExtractUsage, error) {
	contentType := req.ContentType
	if contentType == "" {
		contentType = "image/png"
	}
	imageURL := "data:" + contentType + ";base64," + base64.StdEncoding.EncodeToString(req.Data)
	payload := map[string]any{
		"model": e.model,
		"messages": []map[string]any{
			{
				"role": "user",
				"content": []map[string]any{
					{"type": "text", "text": e.prompt},
					{"type": "image_url", "image_url": map[string]any{"url": imageURL}},
				},
			},
		},
		"temperature": 0,
		"max_tokens":  e.maxTokens,
	}
	body, err := json.Marshal(payload)
	if err != nil {
		return "", ImageExtractUsage{}, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, e.endpoint, bytes.NewReader(body))
	if err != nil {
		return "", ImageExtractUsage{}, err
	}
	httpReq.Header.Set("Authorization", "Bearer "+e.apiKey)
	httpReq.Header.Set("Content-Type", "application/json")

	resp, err := e.client.Do(httpReq)
	if err != nil {
		return "", ImageExtractUsage{}, imageProviderTransportError(err)
	}
	defer func() { _ = resp.Body.Close() }()

	raw, err := io.ReadAll(io.LimitReader(resp.Body, 1<<20))
	if err != nil {
		return "", ImageExtractUsage{}, imageProviderTransportError(err)
	}
	if resp.StatusCode >= 300 {
		return "", ImageExtractUsage{}, imageProviderHTTPError(resp.StatusCode)
	}
	var parsed struct {
		Choices []struct {
			Message struct {
				Content any `json:"content"`
			} `json:"message"`
		} `json:"choices"`
		Usage *struct {
			PromptTokens     int `json:"prompt_tokens"`
			CompletionTokens int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return "", ImageExtractUsage{}, invalidImageProviderResponse(resp.StatusCode, err)
	}
	if len(parsed.Choices) == 0 {
		return "", ImageExtractUsage{}, invalidImageProviderResponse(resp.StatusCode, nil)
	}
	text := extractOpenAIContentText(parsed.Choices[0].Message.Content)
	if strings.TrimSpace(text) == "" {
		return "", ImageExtractUsage{}, invalidImageProviderResponse(resp.StatusCode, nil)
	}
	var usage ImageExtractUsage
	if parsed.Usage != nil {
		usage.PromptTokens = parsed.Usage.PromptTokens
		usage.CompletionTokens = parsed.Usage.CompletionTokens
	}
	return text, usage, nil
}

func extractOpenAIContentText(content any) string {
	switch v := content.(type) {
	case string:
		return strings.TrimSpace(v)
	case []any:
		parts := make([]string, 0, len(v))
		for _, it := range v {
			m, ok := it.(map[string]any)
			if !ok {
				continue
			}
			txt, _ := m["text"].(string)
			if strings.TrimSpace(txt) != "" {
				parts = append(parts, strings.TrimSpace(txt))
			}
		}
		return strings.TrimSpace(strings.Join(parts, "\n"))
	default:
		return ""
	}
}

func truncateString(s string, max int) string {
	if max <= 0 || len(s) <= max {
		return s
	}
	return s[:max]
}