// Package playgroundapp provides endpoints for the model playground.
package playgroundapp

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/ardanlabs/kronk/cmd/server/app/sdk/errs"
	"github.com/ardanlabs/kronk/cmd/server/foundation/logger"
	"github.com/ardanlabs/kronk/cmd/server/foundation/web"
	"github.com/ardanlabs/kronk/sdk/kronk"
	"github.com/ardanlabs/kronk/sdk/kronk/model"
	"github.com/ardanlabs/kronk/sdk/pool"
	"github.com/ardanlabs/kronk/sdk/tools/models"
)

type sessionEntry struct {
	cacheKey      string
	custom        bool
	pendingDelete bool
}

type app struct {
	log      *logger.Logger
	pool     *pool.Pool
	models   *models.Models
	mu       sync.Mutex
	sessions map[string]sessionEntry // session_id -> session entry
	done     chan struct{}
}

func newApp(cfg Config) *app {
	a := &app{
		log:      cfg.Log,
		pool:     cfg.Pool,
		models:   cfg.Models,
		sessions: make(map[string]sessionEntry),
		done:     make(chan struct{}),
	}

	go a.cleanupLoop()

	return a
}

func (a *app) createSession(ctx context.Context, r *http.Request) web.Encoder {
	var req SessionRequest
	if err := web.Decode(r, &req); err != nil {
		return errs.New(errs.InvalidArgument, err)
	}

	if req.ModelID == "" {
		return errs.Errorf(errs.InvalidArgument, "missing model_id")
	}

	baseCfg, err := a.models.KronkResolvedConfig(req.ModelID, a.pool.Kronk.ModelConfig())
	if err != nil {
		return errs.New(errs.Internal, fmt.Errorf("resolving model config: %w", err))
	}

	sessionID, err := generateSessionID()
	if err != nil {
		return errs.New(errs.Internal, fmt.Errorf("generating session id: %w", err))
	}

	cfg := req.Config.ApplyTo(baseCfg)

	// Resolve draft model file paths when the user specifies a draft model ID.
	if req.Config.DraftModelID != nil && *req.Config.DraftModelID != "" {
		draftPath, err := a.models.FullPath(*req.Config.DraftModelID)
		if err != nil {
			return errs.New(errs.InvalidArgument, fmt.Errorf("resolving draft model: %w", err))
		}
		if cfg.DraftModel == nil {
			cfg.DraftModel = &model.DraftModelConfig{}
		}
		cfg.DraftModel.ModelFiles = draftPath.ModelFiles

		// Speculative decoding requires single-slot mode.
		cfg.PtrNSeqMax = new(1)
	}

	if cfg.NUBatch() > cfg.NBatch() {
		return errs.Errorf(errs.InvalidArgument, "nubatch (%d) must not exceed nbatch (%d)", cfg.NUBatch(), cfg.NBatch())
	}

	var (
		cacheKey string
		krn      *kronk.Kronk
	)

	switch {
	case req.HasOverrides():
		cacheKey = fmt.Sprintf("%s/playground/%s", req.ModelID, sessionID)
		krn, err = a.pool.Kronk.AquireCustom(ctx, cacheKey, cfg)
	default:
		cacheKey = req.ModelID
		krn, err = a.pool.Kronk.AquireModel(ctx, req.ModelID)
	}
	if err != nil {
		return errs.New(errs.Internal, err)
	}

	a.log.Info(ctx, "playground-session",
		"session-id", sessionID,
		"model-id", req.ModelID,
		"cache-key", cacheKey,
		"shared", !req.HasOverrides(),
		"krn-ptr", fmt.Sprintf("%p", krn),
		"context-window", krn.ModelConfig().ContextWindow(),
		"nbatch", krn.ModelConfig().NBatch(),
		"nubatch", krn.ModelConfig().NUBatch(),
		"flash-attention", krn.ModelConfig().FlashAttention.String(),
		"cache-type-k", krn.ModelConfig().CacheTypeK.String(),
		"cache-type-v", krn.ModelConfig().CacheTypeV.String(),
		"nseq-max", krn.ModelConfig().NSeqMax(),
	)

	a.mu.Lock()
	a.sessions[sessionID] = sessionEntry{
		cacheKey: cacheKey,
		custom:   req.HasOverrides(),
	}
	a.mu.Unlock()

	effectiveConfig := map[string]any{
		"context_window":    krn.ModelConfig().ContextWindow(),
		"nbatch":            krn.ModelConfig().NBatch(),
		"nubatch":           krn.ModelConfig().NUBatch(),
		"nseq_max":          krn.ModelConfig().NSeqMax(),
		"flash_attention":   krn.ModelConfig().FlashAttention.String(),
		"cache_type_k":      krn.ModelConfig().CacheTypeK.String(),
		"cache_type_v":      krn.ModelConfig().CacheTypeV.String(),
		"incremental_cache": krn.ModelConfig().IncrementalCache(),
		"split_mode":        formatSplitMode(krn.ModelConfig().PtrSplitMode),
		"model_type":        krn.ModelInfo().Type.String(),
		"is_gpt_model":      krn.ModelInfo().IsGPTModel,
	}

	// Report the active drafter. A separate-GGUF draft carries model files;
	// an MTP nDraft override carries only the draft-token count.
	if dm := krn.ModelConfig().DraftModel; dm != nil {
		if dm.IsSeparate() {
			effectiveConfig["draft_model"] = dm.ModelFiles[0]
		}
		effectiveConfig["draft_ndraft"] = dm.NDraft
	}

	return SessionResponse{
		SessionID:       sessionID,
		CacheKey:        cacheKey,
		Status:          "loaded",
		EffectiveConfig: effectiveConfig,
	}
}

func (a *app) deleteSession(ctx context.Context, r *http.Request) web.Encoder {
	id := web.Param(r, "id")
	if id == "" {
		return errs.Errorf(errs.InvalidArgument, "missing session id")
	}

	a.mu.Lock()
	entry, exists := a.sessions[id]
	if !exists {
		a.mu.Unlock()
		return errs.Errorf(errs.NotFound, "session not found: %s", id)
	}

	if !entry.custom {
		delete(a.sessions, id)
		a.mu.Unlock()
		return SessionDeleteResponse{Status: "unloaded"}
	}

	a.mu.Unlock()

	if krn, found := a.pool.Kronk.GetExisting(entry.cacheKey); found && krn.ActiveStreams() > 0 {
		a.mu.Lock()
		entry.pendingDelete = true
		a.sessions[id] = entry
		a.mu.Unlock()
		a.log.Info(ctx, "playground-delete-deferred", "session-id", id, "cache-key", entry.cacheKey, "active-streams", krn.ActiveStreams())
		return SessionDeleteResponse{Status: "unloaded"}
	}

	a.pool.Kronk.Invalidate(entry.cacheKey)

	a.mu.Lock()
	delete(a.sessions, id)
	a.mu.Unlock()

	return SessionDeleteResponse{Status: "unloaded"}
}

func (a *app) chatCompletions(ctx context.Context, r *http.Request) web.Encoder {
	var raw model.D
	if err := json.NewDecoder(r.Body).Decode(&raw); err != nil {
		return errs.New(errs.InvalidArgument, err)
	}

	sessionIDRaw, exists := raw["session_id"]
	if !exists {
		return errs.Errorf(errs.InvalidArgument, "missing session_id field")
	}

	sessionID, ok := sessionIDRaw.(string)
	if !ok {
		return errs.Errorf(errs.InvalidArgument, "session_id must be a string")
	}

	a.mu.Lock()
	entry, exists := a.sessions[sessionID]
	a.mu.Unlock()

	if !exists {
		return errs.Errorf(errs.NotFound, "session not found or expired: %s", sessionID)
	}

	if entry.pendingDelete {
		return errs.Errorf(errs.InvalidArgument, "session is being deleted")
	}

	krn, found := a.pool.Kronk.GetExisting(entry.cacheKey)
	if !found {
		a.mu.Lock()
		delete(a.sessions, sessionID)
		a.mu.Unlock()
		return errs.Errorf(errs.NotFound, "session expired: %s", sessionID)
	}

	ctx, cancel := context.WithTimeout(ctx, 29*time.Minute)
	defer cancel()

	d := model.MapToModelD(raw)

	if _, err := krn.ChatStreamingHTTP(ctx, web.GetWriter(ctx), d); err != nil {

		// Request exceeded the 29-minute deadline. Clean up the session
		// so the next automated-test trial starts with a clean slate.
		// Only invalidate the cache entry for custom (session-scoped)
		// models; shared models (cacheKey == modelID) may be serving
		// other sessions and must not be unloaded.
		if ctx.Err() == context.DeadlineExceeded {
			a.log.Info(ctx, "playground-chat-timeout", "session-id", sessionID, "cache-key", entry.cacheKey)
			if entry.custom {
				a.pool.Kronk.Invalidate(entry.cacheKey)
			}
			a.mu.Lock()
			delete(a.sessions, sessionID)
			a.mu.Unlock()
			return web.NewNoResponse()
		}

		// Client disconnected — streaming headers are already committed
		// so we cannot write an HTTP error response.
		if ctx.Err() != nil {
			return web.NewNoResponse()
		}

		return errs.New(errs.Internal, err)
	}

	return web.NewNoResponse()
}

func generateSessionID() (string, error) {
	b := make([]byte, 16)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return "pg-" + hex.EncodeToString(b), nil
}

func (a *app) cleanupLoop() {
	ticker := time.NewTicker(30 * time.Second)
	defer ticker.Stop()

	for {
		select {
		case <-a.done:
			return
		case <-ticker.C:
			a.mu.Lock()
			var pending []struct {
				id       string
				cacheKey string
			}
			for id, entry := range a.sessions {
				if entry.pendingDelete && entry.custom {
					pending = append(pending, struct {
						id       string
						cacheKey string
					}{id: id, cacheKey: entry.cacheKey})
				}
			}
			a.mu.Unlock()

			for _, p := range pending {
				krn, found := a.pool.Kronk.GetExisting(p.cacheKey)
				if !found || krn.ActiveStreams() == 0 {
					a.pool.Kronk.Invalidate(p.cacheKey)
					a.mu.Lock()
					delete(a.sessions, p.id)
					a.mu.Unlock()
				}
			}

			// Prune sessions whose cache entry has been evicted.
			// Collect candidates under lock, then check cache without holding it.
			type candidate struct {
				id       string
				cacheKey string
			}
			var stale []candidate
			a.mu.Lock()
			for id, entry := range a.sessions {
				if !entry.pendingDelete {
					stale = append(stale, candidate{id: id, cacheKey: entry.cacheKey})
				}
			}
			a.mu.Unlock()

			for _, c := range stale {
				if _, found := a.pool.Kronk.GetExisting(c.cacheKey); !found {
					a.mu.Lock()
					delete(a.sessions, c.id)
					a.mu.Unlock()
				}
			}
		}
	}
}

func formatSplitMode(sm *model.SplitMode) string {
	if sm == nil {
		return "auto"
	}
	return sm.String()
}
