package server

import (
	"ai-router/internal/protocol"
	"bytes"
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"strings"
	"time"
	"unicode/utf8"
)

const modelTestTimeout = 30 * time.Second
const modelTestMaxBody = 1 << 20

type modelTestResult struct {
	OK         bool   `json:"ok"`
	DurationMS int64  `json:"duration_ms"`
	Status     int    `json:"status"`
	Error      string `json:"error,omitempty"`
}

// testModel is an administrative diagnostic, excluded from call logs and usage
// statistics. It creates no API key; upstream providers may still bill the probe.
func (s *Server) testModel(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		errJSON(w, 405, "method not allowed")
		return
	}
	var input struct {
		Model  string  `json:"model"`
		Prompt *string `json:"prompt"`
	}
	if decode(w, r, &input) != nil {
		errJSON(w, 400, "invalid request")
		return
	}
	name, alias, ok := strings.Cut(input.Model, "/")
	prompt := "Reply with OK."
	if input.Prompt != nil {
		prompt = *input.Prompt
	}
	if !ok || name == "" || alias == "" || strings.TrimSpace(prompt) == "" || utf8.RuneCountInString(prompt) > 2000 {
		errJSON(w, 400, "invalid model or prompt")
		return
	}
	ctx, cancel := context.WithTimeout(r.Context(), modelTestTimeout)
	defer cancel()
	var c channel
	var models string
	e := s.db.QueryRowContext(ctx, "SELECT base_url,protocol,enabled,models,api_key FROM channels WHERE name=?", name).Scan(&c.BaseURL, &c.Protocol, &c.Enabled, &models, &c.APIKey)
	if e == sql.ErrNoRows {
		errJSON(w, 403, "model or channel unavailable")
		return
	}
	if e != nil || json.Unmarshal([]byte(models), &c.Models) != nil {
		errJSON(w, 500, "database error")
		return
	}
	actual, mapped := c.Models[alias]
	if !c.Enabled || !mapped || actual == "" {
		errJSON(w, 403, "model or channel unavailable")
		return
	}
	writeJSON(w, s.probeModel(ctx, c, actual, prompt))
}

func (s *Server) probeModel(ctx context.Context, c channel, actual, prompt string) (result modelTestResult) {
	start := time.Now()
	defer func() {
		result.DurationMS = time.Since(start).Milliseconds()
		if ctx.Err() != nil {
			result.OK = false
			result.Error = "model test canceled"
			if ctx.Err() == context.DeadlineExceeded {
				result.Error = "model test timed out"
			}
		}
	}()
	result.Error = "upstream request failed"
	up := protocol.Kind(c.Protocol)
	body, _ := json.Marshal(map[string]any{"messages": []any{map[string]string{"role": "user", "content": prompt}}, "stream": false, "max_completion_tokens": 128})
	payload, e := protocol.ConvertRequest(protocol.Chat, up, body, actual)
	if e != nil {
		result.Error = "invalid upstream configuration"
		return
	}
	suffix := map[protocol.Kind]string{protocol.Chat: "/chat/completions", protocol.Responses: "/responses", protocol.Anthropic: "/v1/messages"}[up]
	req, e := http.NewRequestWithContext(ctx, http.MethodPost, c.BaseURL+suffix, bytes.NewReader(payload))
	if e != nil {
		result.Error = "invalid upstream configuration"
		return
	}
	req.Header.Set("Content-Type", "application/json")
	if up == protocol.Anthropic {
		req.Header.Set("x-api-key", c.APIKey)
		req.Header.Set("anthropic-version", "2023-06-01")
	} else {
		req.Header.Set("Authorization", "Bearer "+c.APIKey)
	}
	client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	res, e := client.Do(req)
	if e != nil {
		if errors.Is(e, context.DeadlineExceeded) {
			result.Error = "model test timed out"
		}
		return
	}
	defer res.Body.Close()
	result.Status = res.StatusCode
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		result.Error = "upstream rejected request"
		return
	}
	body, e = io.ReadAll(io.LimitReader(res.Body, modelTestMaxBody+1))
	if len(body) > modelTestMaxBody {
		result.Error = "upstream response too large"
		return
	}
	if e != nil || !validModelTestResponse(up, body) {
		result.Error = "invalid upstream response"
		return
	}
	result.OK, result.Error = true, ""
	return
}

func validModelTestResponse(kind protocol.Kind, body []byte) bool {
	var r struct {
		ID         string          `json:"id"`
		Object     string          `json:"object"`
		Type       string          `json:"type"`
		Role       string          `json:"role"`
		Error      json.RawMessage `json:"error"`
		Status     string          `json:"status"`
		StopReason string          `json:"stop_reason"`
		Incomplete struct {
			Reason string `json:"reason"`
		} `json:"incomplete_details"`
		Output  []json.RawMessage `json:"output"`
		Content []json.RawMessage `json:"content"`
		Choices []struct {
			Finish  string `json:"finish_reason"`
			Message struct {
				Role    string  `json:"role"`
				Content *string `json:"content"`
			} `json:"message"`
		} `json:"choices"`
	}
	if json.Unmarshal(body, &r) != nil || r.ID == "" || (len(r.Error) != 0 && string(r.Error) != "null") {
		return false
	}
	switch kind {
	case protocol.Chat:
		if r.Object != "chat.completion" || len(r.Choices) == 0 {
			return false
		}
		for _, c := range r.Choices {
			if c.Message.Role != "assistant" || (c.Finish != "stop" && c.Finish != "length") || (c.Message.Content == nil && c.Finish != "length") {
				return false
			}
		}
		return true
	case protocol.Responses:
		if r.Object != "response" || r.Output == nil {
			return false
		}
		exhausted := r.Status == "incomplete" && r.Incomplete.Reason == "max_output_tokens"
		if !exhausted && (r.Status != "completed" || len(r.Output) == 0) {
			return false
		}
		for _, raw := range r.Output {
			var item struct {
				Type    string            `json:"type"`
				Role    string            `json:"role"`
				Content []json.RawMessage `json:"content"`
				Summary []json.RawMessage `json:"summary"`
			}
			if json.Unmarshal(raw, &item) != nil {
				return false
			}
			if item.Type == "reasoning" && item.Summary != nil {
				continue
			}
			if item.Type != "message" || item.Role != "assistant" || !validProbeBlocks(item.Content, "output_text") {
				return false
			}
		}
		return true
	case protocol.Anthropic:
		if r.Type != "message" || r.Role != "assistant" || r.Content == nil {
			return false
		}
		if r.StopReason == "max_tokens" && len(r.Content) == 0 {
			return true
		}
		return (r.StopReason == "end_turn" || r.StopReason == "max_tokens") && validProbeBlocks(r.Content, "text")
	}
	return false
}

func validProbeBlocks(blocks []json.RawMessage, typ string) bool {
	if len(blocks) == 0 {
		return false
	}
	for _, raw := range blocks {
		var b struct {
			Type      string  `json:"type"`
			Text      *string `json:"text"`
			Thinking  *string `json:"thinking"`
			Signature *string `json:"signature"`
			Data      *string `json:"data"`
		}
		if json.Unmarshal(raw, &b) != nil {
			return false
		}
		if typ == "text" && ((b.Type == "thinking" && b.Thinking != nil && b.Signature != nil) || (b.Type == "redacted_thinking" && b.Data != nil)) {
			continue
		}
		if b.Type != typ || b.Text == nil {
			return false
		}
	}
	return true
}