package server

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"io"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"
)

const (
	discoveryTimeout  = 30 * time.Second
	discoveryMaxBody  = 1 << 20
	discoveryMaxPages = 20
)

func (s *Server) discoverModels(w http.ResponseWriter, r *http.Request) {
	if r.Method != http.MethodPost {
		errJSON(w, 405, "method not allowed")
		return
	}
	var c channel
	if decode(w, r, &c) != nil {
		errJSON(w, 400, "invalid channel")
		return
	}
	u, e := url.Parse(c.BaseURL)
	if e != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Hostname() == "" || u.User != nil || u.RawQuery != "" || u.ForceQuery || strings.Contains(c.BaseURL, "#") || (c.Protocol != "openai-chat" && c.Protocol != "openai-responses" && c.Protocol != "anthropic") {
		errJSON(w, 400, "invalid channel")
		return
	}
	if c.ID != "" {
		var old string
		e = s.db.QueryRowContext(r.Context(), "SELECT api_key FROM channels WHERE id=?", c.ID).Scan(&old)
		if e == sql.ErrNoRows {
			errJSON(w, 404, "not found")
			return
		}
		if e != nil {
			errJSON(w, 500, "database error")
			return
		}
		if c.APIKey == "" {
			c.APIKey = old
		}
	}
	ctx, cancel := context.WithTimeout(r.Context(), discoveryTimeout)
	defer cancel()
	models, e := s.fetchModels(ctx, c)
	if e != nil {
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			errJSON(w, 504, "model discovery timed out")
		} else {
			errJSON(w, 502, "model discovery failed")
		}
		return
	}
	writeJSON(w, map[string][]string{"models": models})
}

func (s *Server) fetchModels(ctx context.Context, c channel) ([]string, error) {
	// Share the bounded transport, but never forward credentials across redirects.
	client := &http.Client{Transport: s.client.Transport, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	endpoint := strings.TrimRight(c.BaseURL, "/") + "/models"
	if c.Protocol == "anthropic" {
		endpoint = strings.TrimRight(c.BaseURL, "/") + "/v1/models"
	}
	seen := map[string]bool{}
	cursors := map[string]bool{}
	after := ""
	for page := 0; page < discoveryMaxPages; page++ {
		u := endpoint
		if after != "" {
			u += "?" + url.Values{"after_id": {after}}.Encode()
		}
		req, e := http.NewRequestWithContext(ctx, http.MethodGet, u, nil)
		if e != nil {
			return nil, e
		}
		req.Header.Set("Accept", "application/json")
		if c.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)
		}
		res, e := client.Do(req)
		if e != nil {
			return nil, e
		}
		if res.StatusCode != http.StatusOK {
			res.Body.Close()
			return nil, errors.New("upstream status")
		}
		body, e := io.ReadAll(io.LimitReader(res.Body, discoveryMaxBody+1))
		res.Body.Close()
		if e != nil || len(body) > discoveryMaxBody {
			return nil, errors.New("upstream body")
		}
		var result struct {
			Data []struct {
				ID string `json:"id"`
			} `json:"data"`
			HasMore bool   `json:"has_more"`
			LastID  string `json:"last_id"`
		}
		if json.Unmarshal(body, &result) != nil || result.Data == nil {
			return nil, errors.New("upstream models")
		}
		for _, model := range result.Data {
			if strings.TrimSpace(model.ID) == "" {
				return nil, errors.New("upstream model id")
			}
			seen[model.ID] = true
		}
		if c.Protocol != "anthropic" || !result.HasMore {
			models := make([]string, 0, len(seen))
			for id := range seen {
				models = append(models, id)
			}
			sort.Strings(models)
			return models, nil
		}
		if result.LastID == "" || cursors[result.LastID] {
			return nil, errors.New("upstream pagination")
		}
		cursors[result.LastID] = true
		after = result.LastID
	}
	return nil, errors.New("upstream page limit")
}