package server

import (
	"ai-router/internal/protocol"
	"context"
	"database/sql"
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"net/http"
	"sort"
	"strings"
	"sync"
	"time"
	_ "time/tzdata"

	"modernc.org/sqlite"
)

func init() {
	var locationMu sync.Mutex
	var cachedZone string
	var cachedLocation *time.Location
	if err := sqlite.RegisterDeterministicScalarFunction("router_log_day", 2, func(_ *sqlite.FunctionContext, args []driver.Value) (driver.Value, error) {
		stampText, ok := args[0].(string)
		zone, zoneOK := args[1].(string)
		if !ok || !zoneOK {
			return nil, fmt.Errorf("invalid log timestamp or timezone")
		}
		stamp, err := time.Parse(time.RFC3339Nano, stampText)
		if err != nil {
			return nil, err
		}
		locationMu.Lock()
		defer locationMu.Unlock()
		if cachedLocation == nil || cachedZone != zone {
			loc, err := time.LoadLocation(zone)
			if err != nil {
				return nil, err
			}
			cachedZone, cachedLocation = zone, loc
		}
		return stamp.In(cachedLocation).Format("2006-01-02"), nil
	}); err != nil {
		panic(err)
	}
}

type requestLog struct {
	ID           int64           `json:"id"`
	CreatedAt    string          `json:"created_at"`
	KeyID        string          `json:"key_id"`
	KeyName      string          `json:"key_name"`
	Channel      string          `json:"channel"`
	RequestModel string          `json:"request_model"`
	ActualModel  string          `json:"actual_model"`
	Downstream   string          `json:"downstream"`
	Upstream     string          `json:"upstream"`
	Status       int             `json:"status"`
	DurationMS   int64           `json:"duration_ms"`
	Stream       bool            `json:"stream"`
	Input        *int64          `json:"input_tokens"`
	Output       *int64          `json:"output_tokens"`
	CacheRead    *int64          `json:"cache_read_tokens"`
	CacheWrite   *int64          `json:"cache_write_tokens"`
	UsageKnown   bool            `json:"usage_known"`
	Error        string          `json:"error"`
	RawUsage     json.RawMessage `json:"raw_usage"`
	Usage        protocol.Usage  `json:"-"`
}

func insertLog(ctx context.Context, tx *sql.Tx, l requestLog) error {
	l = boundedLog(l)
	var in, out, cr, cw any
	if l.Usage.Known || len(l.Usage.Raw) > 0 {
		in = l.Usage.Input
		out = l.Usage.Output
		cr = l.Usage.CacheRead
		cw = l.Usage.CacheWrite
	}
	_, err := tx.ExecContext(ctx, "INSERT INTO logs(created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)", l.CreatedAt, l.KeyID, l.KeyName, l.Channel, l.RequestModel, l.ActualModel, l.Downstream, l.Upstream, l.Status, l.DurationMS, l.Stream, in, out, cr, cw, l.Usage.Known, l.Error, string(l.Usage.Raw))
	return err
}
func filter(r *http.Request) (string, []any, error) {
	parts := []string{"1=1"}
	args := []any{}
	for _, name := range []string{"channel", "model"} {
		if v := r.URL.Query().Get(name); v != "" {
			col := name
			if name == "model" {
				col = "actual_model"
			}
			parts = append(parts, col+"=?")
			args = append(args, v)
		}
	}
	if k := r.URL.Query().Get("key_id"); k != "" {
		parts = append(parts, "key_id=?")
		args = append(args, k)
	}
	var from, to time.Time
	for _, n := range []string{"from", "to"} {
		if v := r.URL.Query().Get(n); v != "" {
			t, e := time.Parse(time.RFC3339Nano, v)
			if e != nil {
				return "", nil, e
			}
			if n == "from" {
				from = t
				parts = append(parts, "julianday(created_at)>=julianday(?)")
			} else {
				to = t
				parts = append(parts, "julianday(created_at)<=julianday(?)")
			}
			args = append(args, t.UTC().Format(time.RFC3339Nano))
		}
	}
	if !from.IsZero() && !to.IsZero() && from.After(to) {
		return "", nil, sql.ErrNoRows
	}
	return strings.Join(parts, " AND "), args, nil
}
func (s *Server) readLogs(r *http.Request, limited bool) ([]requestLog, error) {
	where, args, e := filter(r)
	if e != nil {
		return nil, e
	}
	q := "SELECT id,created_at,key_id,key_name,channel,request_model,actual_model,downstream,upstream,status,duration_ms,stream,input_tokens,output_tokens,cache_read_tokens,cache_write_tokens,usage_known,error,raw_usage FROM logs WHERE " + where + " ORDER BY id DESC"
	if limited {
		q += " LIMIT 200"
	}
	rows, e := s.db.Query(q, args...)
	if e != nil {
		return nil, e
	}
	defer rows.Close()
	out := []requestLog{}
	for rows.Next() {
		var l requestLog
		var raw string
		if e = rows.Scan(&l.ID, &l.CreatedAt, &l.KeyID, &l.KeyName, &l.Channel, &l.RequestModel, &l.ActualModel, &l.Downstream, &l.Upstream, &l.Status, &l.DurationMS, &l.Stream, &l.Input, &l.Output, &l.CacheRead, &l.CacheWrite, &l.UsageKnown, &l.Error, &raw); e != nil {
			return nil, e
		}
		if json.Valid([]byte(raw)) {
			l.RawUsage = json.RawMessage(raw)
		}
		out = append(out, l)
	}
	return out, rows.Err()
}
func (s *Server) logs(w http.ResponseWriter, r *http.Request) {
	if _, _, e := filter(r); e != nil {
		errJSON(w, 400, "invalid time filter")
		return
	}
	ls, e := s.readLogs(r, true)
	if e != nil {
		errJSON(w, 500, "database error")
		return
	}
	writeJSON(w, ls)
}

type summary struct {
	Requests   int64  `json:"requests"`
	Success    int64  `json:"success"`
	Failed     int64  `json:"failed"`
	Input      *int64 `json:"input_tokens"`
	Output     *int64 `json:"output_tokens"`
	CacheRead  *int64 `json:"cache_read_tokens"`
	CacheWrite *int64 `json:"cache_write_tokens"`
	Unknown    int64  `json:"unknown_usage"`
}

func (a *summary) add(l requestLog) {
	a.Requests++
	if l.Status >= 200 && l.Status < 300 && l.Error == "" {
		a.Success++
	} else {
		a.Failed++
	}
	if !l.UsageKnown {
		a.Unknown++
		return
	}
	if a.Input == nil {
		a.Input = new(int64)
		a.Output = new(int64)
		a.CacheRead = new(int64)
		a.CacheWrite = new(int64)
	}
	if l.Input != nil {
		*a.Input += *l.Input
	}
	if l.Output != nil {
		*a.Output += *l.Output
	}
	if l.CacheRead != nil {
		*a.CacheRead += *l.CacheRead
	}
	if l.CacheWrite != nil {
		*a.CacheWrite += *l.CacheWrite
	}
}
func (s *Server) stats(w http.ResponseWriter, r *http.Request) {
	zone := r.URL.Query().Get("timezone")
	if zone == "" {
		zone = "UTC"
	}
	_, err := time.LoadLocation(zone)
	if err != nil || zone == "Local" {
		errJSON(w, 400, "invalid timezone")
		return
	}
	where, args, e := filter(r)
	if e != nil {
		errJSON(w, 400, "invalid time filter")
		return
	}
	const aggregates = `COUNT(*), COALESCE(SUM(status>=200 AND status<300 AND error=''),0),
COALESCE(SUM(NOT(status>=200 AND status<300 AND error='')),0),
SUM(CASE WHEN usage_known THEN COALESCE(input_tokens,0) END),
SUM(CASE WHEN usage_known THEN COALESCE(output_tokens,0) END),
SUM(CASE WHEN usage_known THEN COALESCE(cache_read_tokens,0) END),
SUM(CASE WHEN usage_known THEN COALESCE(cache_write_tokens,0) END),
COALESCE(SUM(NOT usage_known),0), MAX(id)`
	// NOT MATERIALIZED avoids copying the full log payload into a temporary table.
	query := `WITH filtered AS NOT MATERIALIZED (SELECT * FROM logs WHERE ` + where + `)
SELECT 'total','','',` + aggregates + ` FROM filtered
UNION ALL SELECT 'key',key_id,key_name,` + aggregates + ` FROM filtered GROUP BY key_id
UNION ALL SELECT 'day',router_log_day(created_at,?), '',` + aggregates + ` FROM filtered GROUP BY router_log_day(created_at,?)
UNION ALL SELECT 'model',channel,actual_model,` + aggregates + ` FROM filtered GROUP BY channel,actual_model`
	args = append(args, zone, zone)
	ctx, cancel := context.WithTimeout(r.Context(), 10*time.Second)
	defer cancel()
	rows, e := s.db.QueryContext(ctx, query, args...)
	if e != nil {
		errJSON(w, 500, "database error")
		return
	}
	defer rows.Close()
	type byKey struct {
		KeyID string `json:"key_id"`
		Name  string `json:"name"`
		summary
	}
	type daily struct {
		Day string `json:"day"`
		summary
	}
	type byModel struct {
		Channel string `json:"channel"`
		Model   string `json:"model"`
		summary
	}
	total := summary{}
	ks := map[string]*byKey{}
	ds := map[string]*daily{}
	ms := map[[2]string]*byModel{}
	groups := 0
	for rows.Next() {
		groups++
		if groups > 10000 {
			errJSON(w, 422, "too many statistics groups; narrow the time range")
			return
		}
		var kind, a, b string
		var v summary
		var latest sql.NullInt64
		if err := rows.Scan(&kind, &a, &b, &v.Requests, &v.Success, &v.Failed, &v.Input, &v.Output, &v.CacheRead, &v.CacheWrite, &v.Unknown, &latest); err != nil {
			errJSON(w, 500, "database error")
			return
		}
		switch kind {
		case "total":
			total = v
		case "key":
			// SQLite's sole MAX selects key_name from the latest ID, not timestamp.
			ks[a] = &byKey{KeyID: a, Name: b, summary: v}
		case "day":
			ds[a] = &daily{Day: a, summary: v}
		case "model":
			ms[[2]string{a, b}] = &byModel{Channel: a, Model: b, summary: v}
		}
	}
	if err := rows.Err(); err != nil {
		errJSON(w, 500, "database error")
		return
	}
	keys := []*byKey{}
	days := []*daily{}
	mods := []*byModel{}
	for _, v := range ks {
		keys = append(keys, v)
	}
	for _, v := range ds {
		days = append(days, v)
	}
	for _, v := range ms {
		mods = append(mods, v)
	}
	sort.Slice(keys, func(i, j int) bool { return keys[i].KeyID < keys[j].KeyID })
	sort.Slice(days, func(i, j int) bool { return days[i].Day < days[j].Day })
	sort.Slice(mods, func(i, j int) bool { return mods[i].Channel+mods[i].Model < mods[j].Channel+mods[j].Model })
	writeJSON(w, map[string]any{"summary": total, "by_key": keys, "daily": days, "by_model": mods, "timezone": zone})
}