package server

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"crypto/subtle"
	"database/sql"
	"encoding/hex"
	"encoding/json"
	"errors"
	"io"
	"io/fs"
	_ "modernc.org/sqlite"
	"net"
	"net/http"
	"net/netip"
	"net/url"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"sync/atomic"
	"time"
)

type Server struct {
	logRetentionDays atomic.Int64
	trustedProxies   []netip.Prefix
	limiter          *requestLimiter
	db               *sql.DB
	password         [32]byte
	assets           fs.FS
	mu               sync.Mutex
	sessions         map[string]time.Time
	attempts         map[string][]time.Time
	authFailures     map[string][]time.Time
	authFailuresSeen time.Time
	loginFailures    []time.Time
	client           *http.Client
	publicOrigin     string
	lifecycleMu      sync.Mutex
	closing          bool
	active           sync.WaitGroup
	closeOnce        sync.Once
	closeErr         error
	logMu            sync.RWMutex
	logsClosed       bool
	logQueue         chan logEntry
	logDone          chan struct{}
	logErr           error
}

const maxBody = 16 << 20

func New(p, password string, a fs.FS) (*Server, error) {
	if password == "" {
		return nil, errors.New("empty admin password")
	}
	if p != ":memory:" {
		if e := os.MkdirAll(filepath.Dir(p), 0700); e != nil {
			return nil, e
		}
	}
	d, e := sql.Open("sqlite", p)
	if e != nil {
		return nil, e
	}
	d.SetMaxOpenConns(1)
	schema := `PRAGMA busy_timeout=5000; CREATE TABLE IF NOT EXISTS channels(id TEXT PRIMARY KEY,name TEXT NOT NULL UNIQUE,base_url TEXT NOT NULL,protocol TEXT NOT NULL,enabled INTEGER NOT NULL,passthrough INTEGER NOT NULL,models TEXT NOT NULL,api_key TEXT NOT NULL); CREATE TABLE IF NOT EXISTS keys(id TEXT PRIMARY KEY,name TEXT NOT NULL,prefix TEXT NOT NULL,hash TEXT NOT NULL UNIQUE,enabled INTEGER NOT NULL,all_channels INTEGER NOT NULL,channels TEXT NOT NULL,expires_at TEXT NOT NULL,created_at TEXT NOT NULL,revoked INTEGER NOT NULL DEFAULT 0); CREATE TABLE IF NOT EXISTS logs(id INTEGER PRIMARY KEY AUTOINCREMENT,created_at TEXT NOT NULL,key_id TEXT NOT NULL,key_name TEXT NOT NULL,channel TEXT NOT NULL,request_model TEXT NOT NULL,actual_model TEXT NOT NULL,downstream TEXT NOT NULL,upstream TEXT NOT NULL,status INTEGER NOT NULL,duration_ms INTEGER NOT NULL,stream INTEGER NOT NULL,input_tokens INTEGER,output_tokens INTEGER,cache_read_tokens INTEGER,cache_write_tokens INTEGER,usage_known INTEGER NOT NULL,error TEXT NOT NULL,raw_usage TEXT NOT NULL); CREATE INDEX IF NOT EXISTS logs_time ON logs(created_at); CREATE INDEX IF NOT EXISTS logs_key ON logs(key_id);`
	if _, e = d.Exec(schema); e != nil {
		d.Close()
		return nil, e
	}
	if p != ":memory:" {
		if e = os.Chmod(p, 0600); e != nil {
			d.Close()
			return nil, e
		}
	}
	tr := http.DefaultTransport.(*http.Transport).Clone()
	// Upstreams must not be able to turn compressed streams into proxy-side
	// decompression work; the router asks for identity encoding.
	tr.DisableCompression = true
	tr.ResponseHeaderTimeout = 60 * time.Second
	tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
		c, e := (&net.Dialer{Timeout: 15 * time.Second, KeepAlive: 30 * time.Second}).DialContext(ctx, network, addr)
		if e != nil {
			return nil, e
		}
		return &idleConn{Conn: c}, nil
	}
	s := &Server{db: d, password: sha256.Sum256([]byte(password)), assets: a, sessions: map[string]time.Time{}, attempts: map[string][]time.Time{}, authFailures: map[string][]time.Time{}, client: &http.Client{Transport: tr, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}, logQueue: make(chan logEntry, 1024), logDone: make(chan struct{})}
	go s.writeLogs()
	s.limiter = newRequestLimiter(120, 8, 64)
	return s, nil
}
func (s *Server) Close() error {
	s.closeOnce.Do(func() {
		s.lifecycleMu.Lock()
		s.closing = true
		s.lifecycleMu.Unlock()
		s.active.Wait()
		s.logMu.Lock()
		s.logsClosed = true
		close(s.logQueue)
		s.logMu.Unlock()
		<-s.logDone
		s.client.CloseIdleConnections()
		s.closeErr = errors.Join(s.logErr, s.db.Close())
	})
	return s.closeErr
}
func writeJSON(w http.ResponseWriter, v any) {
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(v)
}
func errJSON(w http.ResponseWriter, n int, m string) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(n)
	json.NewEncoder(w).Encode(map[string]string{"error": m})
}
func decode(w http.ResponseWriter, r *http.Request, v any) error {
	limit := int64(maxBody)
	if r.URL.Path == "/admin/api/login" {
		limit = 16 << 10
	}
	return readBody(w, r, limit, bodyReadTimeout, func() error {
		d := json.NewDecoder(r.Body)
		if e := d.Decode(v); e != nil {
			return e
		}
		var extra any
		if d.Decode(&extra) != io.EOF {
			return errors.New("trailing JSON")
		}
		return nil
	})
}
func randomID() string {
	b := make([]byte, 32)
	if _, e := rand.Read(b); e != nil {
		panic(e)
	}
	return hex.EncodeToString(b)
}
func hash(v string) string { h := sha256.Sum256([]byte(v)); return hex.EncodeToString(h[:]) }
func (s *Server) session(r *http.Request) bool {
	s.mu.Lock()
	defer s.mu.Unlock()
	c, e := r.Cookie("ai_router_session")
	if e != nil {
		return false
	}
	t, ok := s.sessions[c.Value]
	return ok && time.Now().Before(t)
}
func (s *Server) guard(r *http.Request) bool {
	if r.Header.Get("X-Requested-With") != "ai-router" {
		return false
	}
	o, e := url.Parse(r.Header.Get("Origin"))
	if s.publicOrigin != "" {
		// Origins are scheme/host case-insensitive; a mixed-case deployment
		// value must not lock the admin out.
		return e == nil && strings.EqualFold(r.Header.Get("Origin"), s.publicOrigin)
	}
	scheme := "http"
	if r.TLS != nil {
		scheme = "https"
	}
	return e == nil && o.Scheme == scheme && strings.EqualFold(o.Host, r.Host) && o.User == nil && o.RawQuery == "" && o.Fragment == "" && o.Path == ""
}
func (s *Server) secureCookie() bool {
	return len(s.publicOrigin) >= 8 && strings.EqualFold(s.publicOrigin[:8], "https://")
}
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	w = &deadlineWriter{ResponseWriter: w}
	_ = http.NewResponseController(w).SetReadDeadline(time.Now().Add(bodyReadTimeout))
	s.lifecycleMu.Lock()
	if s.closing {
		s.lifecycleMu.Unlock()
		errJSON(w, http.StatusServiceUnavailable, "server shutting down")
		return
	}
	s.active.Add(1)
	s.lifecycleMu.Unlock()
	defer s.active.Done()
	w.Header().Set("X-Content-Type-Options", "nosniff")
	w.Header().Set("Referrer-Policy", "no-referrer")
	w.Header().Set("X-Frame-Options", "DENY")
	w.Header().Set("Content-Security-Policy", "default-src 'none'; base-uri 'none'; frame-ancestors 'none'; form-action 'none'; connect-src 'self'; script-src 'self'; style-src 'self' 'unsafe-inline'; img-src 'self' data:")
	if strings.HasPrefix(r.URL.Path, "/admin/api/") || r.URL.Path == "/admin/api" {
		w.Header().Set("Cache-Control", "no-store")
		s.admin(w, r)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/v1/") {
		s.proxy(w, r)
		return
	}
	if strings.HasPrefix(r.URL.Path, "/api/") {
		http.NotFound(w, r)
		return
	}
	if s.assets == nil || (r.Method != "GET" && r.Method != "HEAD") {
		http.NotFound(w, r)
		return
	}
	p := strings.TrimPrefix(r.URL.Path, "/")
	if p == "" {
		p = "index.html"
	}
	if st, e := fs.Stat(s.assets, p); e != nil || st.IsDir() {
		p = "index.html"
	}
	http.ServeFileFS(w, r, s.assets, p)
}
func (s *Server) admin(w http.ResponseWriter, r *http.Request) {
	if r.Method != "GET" && !s.guard(r) {
		errJSON(w, 403, "forbidden")
		return
	}
	if r.URL.Path == "/admin/api/login" && r.Method == "POST" {
		s.login(w, r)
		return
	}
	if r.URL.Path == "/admin/api/session" && r.Method == "GET" {
		writeJSON(w, map[string]bool{"authenticated": s.session(r)})
		return
	}
	if !s.session(r) {
		errJSON(w, 401, "unauthorized")
		return
	}
	if r.URL.Path == "/admin/api/logout" && r.Method == "POST" {
		c, _ := r.Cookie("ai_router_session")
		s.mu.Lock()
		delete(s.sessions, c.Value)
		s.mu.Unlock()
		http.SetCookie(w, &http.Cookie{Name: "ai_router_session", MaxAge: -1, Path: "/", HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
		writeJSON(w, map[string]bool{"authenticated": false})
		return
	}
	switch {
	case r.URL.Path == "/admin/api/models/test":
		s.testModel(w, r)
	case r.URL.Path == "/admin/api/channels":
		s.channels(w, r)
	case r.URL.Path == "/admin/api/channels/discover-models":
		s.discoverModels(w, r)
	case strings.HasPrefix(r.URL.Path, "/admin/api/channels/"):
		s.channelID(w, r)
	case r.URL.Path == "/admin/api/keys":
		s.keys(w, r)
	case strings.HasPrefix(r.URL.Path, "/admin/api/keys/"):
		s.keyID(w, r)
	case r.URL.Path == "/admin/api/logs" && r.Method == "GET":
		s.logs(w, r)
	case r.URL.Path == "/admin/api/stats" && r.Method == "GET":
		s.stats(w, r)
	default:
		errJSON(w, 404, "not found")
	}
}
// login gates on failure counts BEFORE verifying the password. While an IP is
// locked — or the global failure budget is spent — every attempt, including a
// correct password, gets the same 429 without a comparison, so responses never
// reveal whether a candidate password is right and brute force stays capped
// at loginFailuresPerMinute per IP and loginGlobalFailuresPerMinute overall.
// Successful logins are not failures: their gate slot is released by clearing
// the IP bucket, and only wrong passwords feed the global window.
func (s *Server) login(w http.ResponseWriter, r *http.Request) {
	ip := s.clientIP(r)
	now := time.Now()
	var x struct {
		Password string `json:"password"`
	}
	if decode(w, r, &x) != nil {
		errJSON(w, 400, "invalid request")
		return
	}
	s.mu.Lock()
	pruneFailures(s.attempts, now)
	s.loginFailures = trimRecent(s.loginFailures, now)
	recent := []time.Time{}
	for _, t := range s.attempts[ip] {
		if now.Sub(t) < time.Minute {
			recent = append(recent, t)
		}
	}
	blocked := len(recent) >= loginFailuresPerMinute || len(s.loginFailures) >= loginGlobalFailuresPerMinute || (len(recent) == 0 && len(s.attempts) >= 10000)
	if !blocked {
		// Reserve the gate slot atomically with the check so a concurrent
		// burst cannot slip past the limit; a success erases the bucket below.
		s.attempts[ip] = append(recent, now)
	}
	s.mu.Unlock()
	if blocked {
		w.Header().Set("Retry-After", "60")
		errJSON(w, 429, "too many attempts")
		return
	}
	h := sha256.Sum256([]byte(x.Password))
	if subtle.ConstantTimeCompare(h[:], s.password[:]) != 1 {
		s.mu.Lock()
		s.loginFailures = append(s.loginFailures, now)
		s.mu.Unlock()
		errJSON(w, 401, "invalid credentials")
		return
	}
	v := randomID()
	s.mu.Lock()
	delete(s.attempts, ip)
	for k, t := range s.sessions {
		if now.After(t) {
			delete(s.sessions, k)
		}
	}
	if c, e := r.Cookie("ai_router_session"); e == nil {
		delete(s.sessions, c.Value)
	}
	s.sessions[v] = now.Add(24 * time.Hour)
	s.mu.Unlock()
	http.SetCookie(w, &http.Cookie{Name: "ai_router_session", Value: v, Path: "/", MaxAge: 86400, HttpOnly: true, SameSite: http.SameSiteStrictMode, Secure: r.TLS != nil || s.secureCookie()})
	writeJSON(w, map[string]bool{"authenticated": true})
}