package main

import (
	"context"
	"embed"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/google/uuid"
	"github.com/gorilla/websocket"
	"github.com/orchard9/peach/pkg/persona"
	"github.com/orchard9/peach/pkg/personagen"
	"github.com/orchard9/peach/pkg/personagen/storage"
)

//go:embed index.html
var staticFiles embed.FS

const (
	agentsDir = "agents"
	mediaDir  = "media"
)

var (
	apiKey   = os.Getenv("LAOZHANG_API_KEY")
	upgrader = websocket.Upgrader{
		CheckOrigin: func(r *http.Request) bool { return true },
	}
)

// SafeConn wraps a WebSocket connection with a mutex for thread-safe writes
type SafeConn struct {
	conn *websocket.Conn
	mu   sync.Mutex
}

func (sc *SafeConn) WriteJSON(v interface{}) error {
	sc.mu.Lock()
	defer sc.mu.Unlock()
	return sc.conn.WriteJSON(v)
}

// Hub manages all WebSocket connections for broadcasting
type Hub struct {
	mu      sync.RWMutex
	clients map[*SafeConn]bool
}

func NewHub() *Hub {
	return &Hub{
		clients: make(map[*SafeConn]bool),
	}
}

func (h *Hub) Register(sc *SafeConn) {
	h.mu.Lock()
	h.clients[sc] = true
	h.mu.Unlock()
}

func (h *Hub) Unregister(sc *SafeConn) {
	h.mu.Lock()
	delete(h.clients, sc)
	h.mu.Unlock()
}

func (h *Hub) Broadcast(msg interface{}) {
	h.mu.RLock()
	defer h.mu.RUnlock()
	for sc := range h.clients {
		sc.WriteJSON(msg)
	}
}

var hub = NewHub()

// Agent represents a created agent (simplified view for UI)
type Agent struct {
	ID          string    `json:"id"`
	Name        string    `json:"name"`
	Handle      string    `json:"handle"`
	Gender      string    `json:"gender"`
	Description string    `json:"description"`
	Tags        []string  `json:"tags"`
	AvatarURL   string    `json:"avatar_url,omitempty"`
	BannerURL   string    `json:"banner_url,omitempty"`
	Videos      []string  `json:"videos,omitempty"`
	Images      []string  `json:"images,omitempty"`
	CreatedAt   time.Time `json:"created_at"`
	DirName     string    `json:"dir_name,omitempty"`
}

// MediaJob tracks background media generation
type MediaJob struct {
	ID        string `json:"id"`
	AgentID   string `json:"agent_id"`
	AgentName string `json:"agent_name"`
	Type      string `json:"type"`   // "anchor", "avatar", "banner", "video", "gallery"
	Status    string `json:"status"` // "pending", "generating", "complete", "error"
	Progress  int    `json:"progress"`
	Total     int    `json:"total"`
	Error     string `json:"error,omitempty"`
}

var (
	jobsMu   sync.RWMutex
	jobs     = make(map[string]*MediaJob)
	agentsMu sync.Mutex // protects agent modifications and saves
)

// WSMessage represents a WebSocket message
type WSMessage struct {
	Type    string      `json:"type"`
	Payload interface{} `json:"payload"`
}

// CreateRequest is the payload for creating an agent
type CreateRequest struct {
	Description string `json:"description"`
	Gender      string `json:"gender"`
}

// GenerateMediaRequest is for generating media on existing agent
type GenerateMediaRequest struct {
	AgentID   string `json:"agent_id"`
	MediaType string `json:"media_type"`
}

// DeleteImageRequest is for deleting an image from an agent
type DeleteImageRequest struct {
	AgentID  string `json:"agent_id"`
	ImageURL string `json:"image_url"`
}

// DeleteVideoRequest is for deleting a video from an agent
type DeleteVideoRequest struct {
	AgentID  string `json:"agent_id"`
	VideoURL string `json:"video_url"`
}

// BulkDeleteRequest is for deleting multiple images/videos at once
type BulkDeleteRequest struct {
	AgentID   string   `json:"agent_id"`
	ImageURLs []string `json:"image_urls,omitempty"`
	VideoURLs []string `json:"video_urls,omitempty"`
}

// RegeneratePositionsRequest is for regenerating specific image positions
type RegeneratePositionsRequest struct {
	AgentID   string `json:"agent_id"`
	Positions []int  `json:"positions"`
}

// ModifyPersonaRequest is for modifying persona spec attributes
type ModifyPersonaRequest struct {
	AgentID      string `json:"agent_id"`
	Modification string `json:"modification"` // "tone_eye_color", "align_hair_brows"
	Intensity    string `json:"intensity,omitempty"`
}

// RegenerateBannerRequest is for regenerating banner with optional style
type RegenerateBannerRequest struct {
	AgentID string `json:"agent_id"`
	Style   string `json:"style,omitempty"` // "lifestyle", "portrait", "scenic", "artistic"
}

// BulkRegenerateVideosRequest is for regenerating all videos for multiple agents
type BulkRegenerateVideosRequest struct {
	AgentIDs []string `json:"agent_ids"`
}

func main() {
	if apiKey == "" {
		log.Fatal("LAOZHANG_API_KEY environment variable is required")
	}

	os.MkdirAll(agentsDir, 0755)
	os.MkdirAll(mediaDir, 0755)

	http.HandleFunc("/", serveIndex)
	http.HandleFunc("/ws", handleWebSocket)
	http.HandleFunc("/agents", listAgents)
	http.HandleFunc("/agents/", getAgent)
	http.Handle("/media/", http.StripPrefix("/media/", http.FileServer(http.Dir(mediaDir))))

	port := os.Getenv("PORT")
	if port == "" {
		port = "8080"
	}

	log.Printf("Starting agent creator on http://localhost:%s", port)
	log.Printf("Agents saved to ./%s/, media saved to ./%s/", agentsDir, mediaDir)
	log.Fatal(http.ListenAndServe(":"+port, nil))
}

func serveIndex(w http.ResponseWriter, r *http.Request) {
	if r.URL.Path != "/" {
		http.NotFound(w, r)
		return
	}
	content, _ := staticFiles.ReadFile("index.html")
	w.Header().Set("Content-Type", "text/html")
	w.Write(content)
}

func handleWebSocket(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Printf("WebSocket upgrade failed: %v", err)
		return
	}
	defer conn.Close()

	sc := &SafeConn{conn: conn}
	hub.Register(sc)
	defer hub.Unregister(sc)

	log.Println("WebSocket client connected")

	// Send current jobs
	jobsMu.RLock()
	for _, job := range jobs {
		sc.WriteJSON(WSMessage{Type: "job_update", Payload: job})
	}
	jobsMu.RUnlock()

	for {
		_, message, err := conn.ReadMessage()
		if err != nil {
			if websocket.IsCloseError(err, websocket.CloseNormalClosure, websocket.CloseGoingAway) {
				log.Println("WebSocket client disconnected")
			}
			return
		}

		var msg WSMessage
		if err := json.Unmarshal(message, &msg); err != nil {
			continue
		}

		switch msg.Type {
		case "create":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req CreateRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid create request")
				continue
			}
			go handleCreate(sc, req)
		case "generate_media":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req GenerateMediaRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid generate_media request")
				continue
			}
			go handleGenerateMedia(sc, req)
		case "delete_image":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req DeleteImageRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid delete_image request")
				continue
			}
			handleDeleteImage(sc, req)
		case "delete_video":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req DeleteVideoRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid delete_video request")
				continue
			}
			handleDeleteVideo(sc, req)
		case "bulk_delete":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req BulkDeleteRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid bulk_delete request")
				continue
			}
			handleBulkDelete(sc, req)
		case "clear_jobs":
			clearCompletedJobs()
		case "regenerate_positions":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req RegeneratePositionsRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid regenerate_positions request")
				continue
			}
			go handleRegeneratePositions(sc, req)
		case "modify_persona":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req ModifyPersonaRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid modify_persona request")
				continue
			}
			handleModifyPersona(sc, req)
		case "regenerate_banner":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req RegenerateBannerRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid regenerate_banner request")
				continue
			}
			go handleRegenerateBanner(sc, req)
		case "bulk_regenerate_videos":
			payloadBytes, _ := json.Marshal(msg.Payload)
			var req BulkRegenerateVideosRequest
			if err := json.Unmarshal(payloadBytes, &req); err != nil {
				sendError(sc, "Invalid bulk_regenerate_videos request")
				continue
			}
			go handleBulkRegenerateVideos(sc, req)
		default:
			sendError(sc, "Unknown message type: "+msg.Type)
		}
	}
}

func handleCreate(sc *SafeConn, req CreateRequest) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	send := func(msg WSMessage) {
		sc.WriteJSON(msg)
	}

	// Create personagen service
	cfg := personagen.ServiceConfig{
		LaozhangAPIKey: apiKey,
	}

	progress := &wsProgressReporter{send: send}
	service, err := personagen.NewService(ctx, cfg, progress)
	if err != nil {
		sendError(sc, "Failed to create service: "+err.Error())
		return
	}

	// Step 1: Generate persona spec
	send(WSMessage{Type: "step", Payload: map[string]string{"step": "personality"}})

	// Extract ethnicity from description if present
	ethnicity := personagen.InferEthnicityFromDescription(req.Description)

	seed := personagen.SeedParams{
		Gender:          req.Gender,
		CustomArchetype: req.Description,
		Ethnicity:       ethnicity,
	}

	spec, err := service.Generate(ctx, seed)
	if err != nil {
		sendError(sc, "Failed to generate persona: "+err.Error())
		return
	}
	send(WSMessage{Type: "step_done", Payload: map[string]string{"step": "personality"}})
	send(WSMessage{Type: "step_done", Payload: map[string]string{"step": "appearance"}})
	send(WSMessage{Type: "step_done", Payload: map[string]string{"step": "identity"}})
	send(WSMessage{Type: "step_done", Payload: map[string]string{"step": "tags"}})

	// Validate name length (max 30 chars for UI display)
	fullName := fmt.Sprintf("%s %s", spec.CoreIdentity.Name.First, spec.CoreIdentity.Name.Last)
	if len(fullName) > 30 {
		sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]interface{}{
			"code":    "NAME_TOO_LONG",
			"message": fmt.Sprintf("Generated name too long (%d chars, max 30): %s", len(fullName), fullName),
		}})
		return
	}

	// Check for duplicate agent
	if existing, found := checkDuplicateAgent(fullName); found {
		sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]interface{}{
			"code":        "DUPLICATE_AGENT",
			"message":     fmt.Sprintf("Agent '%s' already exists", fullName),
			"existing_id": existing.ID,
		}})
		return
	}

	// Create storage and save persona
	optimizer := storage.NewOptimizer(85, 80) // WebP quality 85, AVIF quality 80
	store := storage.NewFilesystemStorage(mediaDir, optimizer)

	// Generate persona ID from name
	safeName := strings.ToLower(strings.ReplaceAll(spec.CoreIdentity.Name.First, " ", "_"))
	shortID := uuid.New().String()[:4]
	personaID := fmt.Sprintf("%s_%s", safeName, shortID)

	// Initialize ImageMatrix before saving (so it persists to disk)
	if len(spec.ImageMatrix) == 0 {
		spec.ImageMatrix = persona.ExtendedImageMatrix(75)
	}

	saveResult, err := store.SavePersona(ctx, spec, storage.SaveOptions{
		ID:     personaID,
		Format: storage.FormatJSON,
		Force:  true,
	})
	if err != nil {
		sendError(sc, "Failed to save persona: "+err.Error())
		return
	}

	// Create agent for UI
	agent := &Agent{
		ID:          personaID,
		Name:        fmt.Sprintf("%s %s", spec.CoreIdentity.Name.First, spec.CoreIdentity.Name.Last),
		Handle:      strings.ToLower(spec.CoreIdentity.Name.First),
		Gender:      req.Gender,
		Description: req.Description,
		Tags:        extractTags(spec),
		Images:      []string{},
		CreatedAt:   time.Now(),
		DirName:     personaID,
	}

	saveAgentJSON(agent)

	send(WSMessage{Type: "complete", Payload: map[string]interface{}{"agent": agent}})
	hub.Broadcast(WSMessage{Type: "agent_created", Payload: agent})

	log.Printf("Created agent: %s (%s)", agent.Name, agent.ID)
	log.Printf("Persona saved to: %s", saveResult.Dir)

	// Spawn media generation with its own long-running context
	mediaCtx, mediaCancel := context.WithTimeout(context.Background(), 30*time.Minute)
	go func() {
		defer mediaCancel()
		spawnMediaJobs(mediaCtx, service, store, spec, agent)
	}()
}

func spawnMediaJobs(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	personaID := agent.ID

	// Job 1: Generate anchor image (position 1) - MUST be first
	anchorJob := createJob(agent, "anchor", 1)
	anchorJob.Status = "generating"
	updateJob(anchorJob)

	log.Printf("Generating anchor image for %s...", agent.Name)
	images, err := service.GenerateImages(ctx, spec, []int{1}, nil)
	if err != nil {
		anchorJob.Status = "error"
		anchorJob.Error = err.Error()
		updateJob(anchorJob)
		log.Printf("Anchor generation failed for %s: %v", agent.Name, err)
		return
	}

	anchorBytes := images[1]

	// Save anchor
	anchorPath, err := store.SaveAnchor(ctx, personaID, anchorBytes)
	if err != nil {
		log.Printf("Failed to save anchor: %v", err)
	} else {
		log.Printf("Saved anchor: %s", anchorPath)
	}

	// Save as position 1 image
	if len(spec.ImageMatrix) > 0 {
		paths, err := store.SaveImage(ctx, personaID, 1, anchorBytes, &spec.ImageMatrix[0])
		if err == nil {
			newImagePath := "/media/" + personaID + "/images/" + filepath.Base(paths.WebP)
			updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
				a.Images = append(a.Images, newImagePath)
			})
			if err == nil {
				hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
			}
		}
	}

	// Set anchor for all subsequent generation
	service.SetAnchor(anchorBytes)

	anchorJob.Status = "complete"
	anchorJob.Progress = 100
	updateJob(anchorJob)
	log.Printf("Anchor generated for %s", agent.Name)

	// Now run remaining jobs in parallel
	var wg sync.WaitGroup

	// Job 2: Avatar (uses anchor)
	wg.Add(1)
	go func() {
		defer wg.Done()
		generateAvatarJob(ctx, service, store, spec, agent)
	}()

	// Job 3: Banner (uses anchor)
	wg.Add(1)
	go func() {
		defer wg.Done()
		generateBannerJob(ctx, service, store, spec, agent)
	}()

	// Job 4: Gallery images (uses anchor, runs in parallel batches)
	wg.Add(1)
	go func() {
		defer wg.Done()
		generateGalleryJob(ctx, service, store, spec, agent)
	}()

	// Job 5: Video (uses anchor)
	wg.Add(1)
	go func() {
		defer wg.Done()
		generateVideoJob(ctx, service, store, spec, agent)
	}()

	wg.Wait()
	log.Printf("All media generation complete for %s", agent.Name)
}

func generateAvatarJob(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	job := createJob(agent, "avatar", 1)
	job.Status = "generating"
	updateJob(job)

	result, err := service.GenerateAvatar(ctx, spec)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		log.Printf("Avatar generation failed for %s: %v", agent.Name, err)
		return
	}

	path, err := store.SaveAvatar(ctx, agent.ID, result.Data, result.Prompt)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	avatarURL := fmt.Sprintf("/media/%s/%s?t=%d", agent.ID, filepath.Base(path), time.Now().Unix())
	updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
		a.AvatarURL = avatarURL
	})
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)

	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	log.Printf("Avatar generated for %s: %s", agent.Name, path)
}

func generateBannerJob(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	job := createJob(agent, "banner", 1)
	job.Status = "generating"
	updateJob(job)

	result, err := service.GenerateBanner(ctx, spec)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		log.Printf("Banner generation failed for %s: %v", agent.Name, err)
		return
	}

	path, err := store.SaveBanner(ctx, agent.ID, result.Data, result.Prompt)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	bannerURL := fmt.Sprintf("/media/%s/%s?t=%d", agent.ID, filepath.Base(path), time.Now().Unix())
	updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
		a.BannerURL = bannerURL
	})
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)

	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	log.Printf("Banner generated for %s: %s", agent.Name, path)
}

func generateAnchorJob(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	job := createJob(agent, "anchor", 1)
	job.Status = "generating"
	updateJob(job)

	log.Printf("Regenerating anchor image for %s...", agent.Name)
	images, err := service.GenerateImages(ctx, spec, []int{1}, nil)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		log.Printf("Anchor regeneration failed for %s: %v", agent.Name, err)
		return
	}

	anchorBytes := images[1]

	// Save anchor
	anchorPath, err := store.SaveAnchor(ctx, agent.ID, anchorBytes)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		log.Printf("Failed to save anchor: %v", err)
		return
	}
	log.Printf("Saved anchor: %s", anchorPath)

	// Save as position 1 image (replace existing if any)
	paths, err := store.SaveImage(ctx, agent.ID, 1, anchorBytes, nil)
	if err == nil {
		newImagePath := "/media/" + agent.ID + "/images/" + filepath.Base(paths.WebP)
		updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
			// Replace position 1 if exists, otherwise prepend
			found := false
			for i, imgURL := range a.Images {
				base := filepath.Base(imgURL)
				if len(base) >= 2 {
					var pos int
					fmt.Sscanf(base[:2], "%d", &pos)
					if pos == 1 {
						a.Images[i] = newImagePath
						found = true
						break
					}
				}
			}
			if !found {
				a.Images = append([]string{newImagePath}, a.Images...)
			}
		})
		if err == nil {
			hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
		}
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)
	log.Printf("Anchor regenerated for %s", agent.Name)
}

func generateGalleryJob(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	// Reload agent to get latest state (in case other jobs updated it)
	agentsMu.Lock()
	currentAgent, err := loadAgent(agent.ID)
	agentsMu.Unlock()
	if err != nil {
		log.Printf("Failed to reload agent for gallery: %v", err)
		return
	}

	// Find highest existing position
	highestPos := 1 // position 1 is anchor
	for _, imgURL := range currentAgent.Images {
		base := filepath.Base(imgURL)
		if len(base) >= 2 {
			var pos int
			fmt.Sscanf(base[:2], "%d", &pos)
			if pos > highestPos {
				highestPos = pos
			}
		}
	}

	// Generate next 10 positions (starting after highest, max 75)
	startPos := highestPos + 1
	if startPos > 75 {
		log.Printf("All 75 positions already generated for %s", agent.Name)
		return
	}

	numToGenerate := 10
	if startPos+numToGenerate-1 > 75 {
		numToGenerate = 75 - startPos + 1
	}

	positionsToGenerate := make([]int, numToGenerate)
	for i := 0; i < numToGenerate; i++ {
		positionsToGenerate[i] = startPos + i
	}

	log.Printf("Generating positions %d-%d for %s (have %d images)", startPos, startPos+numToGenerate-1, agent.Name, len(currentAgent.Images))

	job := createJob(agent, "gallery", numToGenerate)
	job.Status = "generating"
	updateJob(job)

	for i, pos := range positionsToGenerate {
		select {
		case <-ctx.Done():
			log.Printf("Gallery generation cancelled for %s", agent.Name)
			job.Status = "error"
			job.Error = "cancelled"
			updateJob(job)
			return
		default:
		}

		log.Printf("Generating gallery image %d/%d for %s (position %d)...", i+1, numToGenerate, agent.Name, pos)

		images, err := service.GenerateImages(ctx, spec, []int{pos}, nil)
		if err != nil {
			log.Printf("Gallery image %d failed for %s: %v", pos, agent.Name, err)
			continue
		}

		imgBytes := images[pos]
		paths, err := store.SaveImage(ctx, agent.ID, pos, imgBytes, nil)
		if err != nil {
			log.Printf("Failed to save gallery image %d: %v", pos, err)
			continue
		}

		newImagePath := "/media/" + agent.ID + "/images/" + filepath.Base(paths.WebP)
		updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
			a.Images = append(a.Images, newImagePath)
		})
		if err == nil {
			hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
		}

		job.Progress = int(float64(i+1) / float64(numToGenerate) * 100)
		updateJob(job)
		log.Printf("Gallery image %d generated for %s", pos, agent.Name)
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)
}

func generateVideoJob(ctx context.Context, service *personagen.Service, store storage.Storage, spec *persona.PersonaSpec, agent *Agent) {
	// Reload agent to get latest state (in case other jobs updated it)
	agentsMu.Lock()
	currentAgent, err := loadAgent(agent.ID)
	agentsMu.Unlock()
	if err != nil {
		log.Printf("Failed to reload agent for video: %v", err)
		return
	}

	// Video positions 1-4, each with a motion type
	motionTypes := []struct {
		position   int
		motionType string
	}{
		{1, storage.MotionSmileReveal},
		{2, storage.MotionPersonalityMoment},
		{3, storage.MotionLifestyle},
		{4, storage.MotionInvitation},
	}

	// Find which videos already exist
	existingVideos := make(map[string]bool)
	for _, url := range currentAgent.Videos {
		existingVideos[filepath.Base(strings.Split(url, "?")[0])] = true
	}

	// Find next video to generate
	var toGenerate *struct {
		position   int
		motionType string
	}
	for _, mt := range motionTypes {
		filename := fmt.Sprintf("%02d_%s.mp4", mt.position, mt.motionType)
		if !existingVideos[filename] {
			toGenerate = &struct {
				position   int
				motionType string
			}{mt.position, mt.motionType}
			break
		}
	}

	if toGenerate == nil {
		log.Printf("All 4 video positions already generated for %s", agent.Name)
		return
	}

	job := createJob(agent, "video", 1)
	job.Status = "generating"
	updateJob(job)

	log.Printf("Generating video position %d (%s) for %s...", toGenerate.position, toGenerate.motionType, agent.Name)

	videos, err := service.GenerateVideos(ctx, spec, []int{toGenerate.position})
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		log.Printf("Video generation failed for %s: %v", agent.Name, err)
		return
	}

	videoBytes := videos[toGenerate.position]
	path, err := store.SaveVideo(ctx, agent.ID, toGenerate.position, toGenerate.motionType, videoBytes)
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	videoURL := fmt.Sprintf("/media/%s/videos/%s?t=%d", agent.ID, filepath.Base(path), time.Now().Unix())

	updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
		a.Videos = append(a.Videos, videoURL)
	})
	if err != nil {
		job.Status = "error"
		job.Error = err.Error()
		updateJob(job)
		return
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)

	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	log.Printf("Video generated for %s: %s", agent.Name, path)
}

func handleGenerateMedia(sc *SafeConn, req GenerateMediaRequest) {
	// Use a single long-running context for the entire operation
	// Don't defer cancel() here - the goroutine will cancel when done
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)

	// Load agent
	agent, err := loadAgent(req.AgentID)
	if err != nil {
		cancel()
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	// Create service with the long-running context
	cfg := personagen.ServiceConfig{
		LaozhangAPIKey: apiKey,
	}
	service, err := personagen.NewService(ctx, cfg, &personagen.NoopProgressReporter{})
	if err != nil {
		cancel()
		sendError(sc, "Failed to create service: "+err.Error())
		return
	}

	// Load persona spec and anchor
	optimizer := storage.NewOptimizer(85, 80)
	store := storage.NewFilesystemStorage(mediaDir, optimizer)
	spec, err := store.LoadPersona(ctx, agent.ID)
	if err != nil {
		cancel()
		sendError(sc, "Failed to load persona: "+err.Error())
		return
	}

	// Ensure ImageMatrix has 75 positions for extended generation
	if len(spec.ImageMatrix) < 75 {
		spec.ImageMatrix = persona.ExtendedImageMatrix(75)
	}

	// Handle anchor regeneration separately (doesn't need existing anchor)
	if req.MediaType == "anchor" {
		go func() {
			defer cancel()
			generateAnchorJob(ctx, service, store, spec, agent)
		}()
		return
	}

	// All other media types require an existing anchor
	anchorBytes, err := store.LoadAnchor(ctx, agent.ID)
	if err != nil {
		cancel()
		// Send specific error so frontend can show "Regenerate Anchor" button
		sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]string{
			"message": "Anchor not found",
			"code":    "ANCHOR_MISSING",
		}})
		return
	}
	service.SetAnchor(anchorBytes)

	switch req.MediaType {
	case "avatar":
		go func() {
			defer cancel()
			generateAvatarJob(ctx, service, store, spec, agent)
		}()
	case "banner":
		go func() {
			defer cancel()
			generateBannerJob(ctx, service, store, spec, agent)
		}()
	case "images":
		go func() {
			defer cancel()
			generateGalleryJob(ctx, service, store, spec, agent)
		}()
	case "video":
		go func() {
			defer cancel()
			generateVideoJob(ctx, service, store, spec, agent)
		}()
	default:
		cancel()
		sendError(sc, "Unknown media type: "+req.MediaType)
	}
}

func handleDeleteImage(sc *SafeConn, req DeleteImageRequest) {
	imageURL := req.ImageURL

	// Delete the actual files (webp, avif, caption)
	// imageURL is like /media/agent_id/images/01_closeup_front_neutral.webp
	relativePath := strings.TrimPrefix(imageURL, "/media/")
	basePath := filepath.Join(mediaDir, relativePath)
	baseWithoutExt := strings.TrimSuffix(basePath, filepath.Ext(basePath))

	// Remove webp, avif, caption files
	os.Remove(basePath)                    // .webp
	os.Remove(baseWithoutExt + ".avif")    // .avif
	os.Remove(baseWithoutExt + ".caption") // .caption

	// Update agent atomically
	var found bool
	updatedAgent, err := updateAgentField(req.AgentID, func(a *Agent) {
		newImages := []string{}
		for _, img := range a.Images {
			if img == imageURL {
				found = true
				continue
			}
			newImages = append(newImages, img)
		}
		a.Images = newImages
	})

	if err != nil {
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	if !found {
		sendError(sc, "Image not found in agent")
		return
	}

	log.Printf("Deleted image %s from %s", imageURL, updatedAgent.Name)
	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	sc.WriteJSON(WSMessage{Type: "image_deleted", Payload: map[string]string{"image_url": imageURL}})
}

func handleDeleteVideo(sc *SafeConn, req DeleteVideoRequest) {
	videoURL := req.VideoURL

	// Strip query params for file path comparison
	urlWithoutQuery := strings.Split(videoURL, "?")[0]

	// Delete the actual video file
	// videoURL is like /media/agent_id/videos/01_smile_reveal.mp4?t=123
	relativePath := strings.TrimPrefix(urlWithoutQuery, "/media/")
	fullPath := filepath.Join(mediaDir, relativePath)

	os.Remove(fullPath)

	// Update agent atomically
	var found bool
	updatedAgent, err := updateAgentField(req.AgentID, func(a *Agent) {
		newVideos := []string{}
		for _, vid := range a.Videos {
			// Strip query params for comparison
			vidWithoutQuery := strings.Split(vid, "?")[0]
			if vidWithoutQuery == urlWithoutQuery {
				found = true
				continue
			}
			newVideos = append(newVideos, vid)
		}
		a.Videos = newVideos
	})

	if err != nil {
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	if !found {
		sendError(sc, "Video not found in agent")
		return
	}

	log.Printf("Deleted video %s from %s", videoURL, updatedAgent.Name)
	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	sc.WriteJSON(WSMessage{Type: "video_deleted", Payload: map[string]string{"video_url": videoURL}})
}

// checkDuplicateAgent checks if an agent with this name already exists
func checkDuplicateAgent(name string) (*Agent, bool) {
	files, _ := os.ReadDir(agentsDir)
	for _, f := range files {
		if !strings.HasSuffix(f.Name(), ".json") {
			continue
		}
		data, _ := os.ReadFile(filepath.Join(agentsDir, f.Name()))
		var agent Agent
		if json.Unmarshal(data, &agent) == nil {
			if strings.EqualFold(agent.Name, name) {
				return &agent, true
			}
		}
	}
	return nil, false
}

func handleRegeneratePositions(sc *SafeConn, req RegeneratePositionsRequest) {
	// Validate positions
	for _, pos := range req.Positions {
		if pos < 1 || pos > 75 {
			sendError(sc, fmt.Sprintf("Invalid position %d: must be between 1 and 75", pos))
			return
		}
	}

	if len(req.Positions) == 0 {
		sendError(sc, "No positions specified")
		return
	}

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

	// Load agent
	agent, err := loadAgent(req.AgentID)
	if err != nil {
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	// Create service
	cfg := personagen.ServiceConfig{
		LaozhangAPIKey: apiKey,
	}
	service, err := personagen.NewService(ctx, cfg, &personagen.NoopProgressReporter{})
	if err != nil {
		sendError(sc, "Failed to create service: "+err.Error())
		return
	}

	// Load persona spec and anchor
	optimizer := storage.NewOptimizer(85, 80)
	store := storage.NewFilesystemStorage(mediaDir, optimizer)
	spec, err := store.LoadPersona(ctx, agent.ID)
	if err != nil {
		sendError(sc, "Failed to load persona: "+err.Error())
		return
	}

	// Ensure ImageMatrix has 75 positions
	if len(spec.ImageMatrix) < 75 {
		spec.ImageMatrix = persona.ExtendedImageMatrix(75)
	}

	// Load anchor for consistency
	anchorBytes, err := store.LoadAnchor(ctx, agent.ID)
	if err != nil {
		sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]string{
			"message": "Anchor not found - regenerate anchor first",
			"code":    "ANCHOR_MISSING",
		}})
		return
	}
	service.SetAnchor(anchorBytes)

	// Notify that regeneration is starting
	sc.WriteJSON(WSMessage{Type: "positions_regenerating", Payload: map[string]interface{}{
		"positions": req.Positions,
		"agent_id":  req.AgentID,
	}})

	// Create job for tracking
	job := createJob(agent, "gallery", len(req.Positions))
	job.Status = "generating"
	updateJob(job)

	log.Printf("Regenerating positions %v for %s...", req.Positions, agent.Name)

	// Delete existing images at these positions first
	for _, pos := range req.Positions {
		deleteImageAtPosition(agent.ID, pos, store)
	}

	// Update agent to remove deleted images
	updatedAgent, _ := updateAgentField(agent.ID, func(a *Agent) {
		newImages := []string{}
		for _, imgURL := range a.Images {
			base := filepath.Base(imgURL)
			keep := true
			for _, pos := range req.Positions {
				prefix := fmt.Sprintf("%02d_", pos)
				if strings.HasPrefix(base, prefix) {
					keep = false
					break
				}
			}
			if keep {
				newImages = append(newImages, imgURL)
			}
		}
		a.Images = newImages
	})
	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})

	// Generate new images for each position
	for i, pos := range req.Positions {
		log.Printf("Regenerating position %d (%d/%d) for %s...", pos, i+1, len(req.Positions), agent.Name)

		images, err := service.GenerateImages(ctx, spec, []int{pos}, nil)
		if err != nil {
			log.Printf("Position %d regeneration failed for %s: %v", pos, agent.Name, err)
			continue
		}

		imgBytes := images[pos]
		var imgMeta *persona.ImageSpec
		if pos <= len(spec.ImageMatrix) {
			imgMeta = &spec.ImageMatrix[pos-1]
		}

		paths, err := store.SaveImage(ctx, agent.ID, pos, imgBytes, imgMeta)
		if err != nil {
			log.Printf("Failed to save regenerated image %d: %v", pos, err)
			continue
		}

		newImagePath := "/media/" + agent.ID + "/images/" + filepath.Base(paths.WebP)
		updatedAgent, err := updateAgentField(agent.ID, func(a *Agent) {
			a.Images = append(a.Images, newImagePath)
		})
		if err == nil {
			hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
		}

		job.Progress = int(float64(i+1) / float64(len(req.Positions)) * 100)
		updateJob(job)
		log.Printf("Position %d regenerated for %s", pos, agent.Name)
	}

	job.Status = "complete"
	job.Progress = 100
	updateJob(job)
	log.Printf("Position regeneration complete for %s", agent.Name)
}

// deleteImageAtPosition removes image files for a specific position
func deleteImageAtPosition(agentID string, position int, store storage.Storage) {
	imagesDir := filepath.Join(mediaDir, agentID, "images")
	prefix := fmt.Sprintf("%02d_", position)

	entries, err := os.ReadDir(imagesDir)
	if err != nil {
		return
	}

	for _, entry := range entries {
		if strings.HasPrefix(entry.Name(), prefix) {
			os.Remove(filepath.Join(imagesDir, entry.Name()))
		}
	}
}

func handleModifyPersona(sc *SafeConn, req ModifyPersonaRequest) {
	ctx := context.Background()

	// Load persona spec
	optimizer := storage.NewOptimizer(85, 80)
	store := storage.NewFilesystemStorage(mediaDir, optimizer)
	spec, err := store.LoadPersona(ctx, req.AgentID)
	if err != nil {
		sendError(sc, "Failed to load persona: "+err.Error())
		return
	}

	var modified bool
	var description string

	switch req.Modification {
	case "tone_eye_color":
		modified = toneDownEyeColor(spec, req.Intensity)
		description = fmt.Sprintf("Eye color toned down (%s)", req.Intensity)

	case "align_hair_brows":
		modified = alignHairAndBrows(spec)
		description = "Brow color aligned to hair color"

	default:
		sendError(sc, "Unknown modification type: "+req.Modification)
		return
	}

	if !modified {
		sc.WriteJSON(WSMessage{Type: "persona_modified", Payload: map[string]interface{}{
			"agent_id":    req.AgentID,
			"modified":    false,
			"description": "No changes needed",
		}})
		return
	}

	// Save modified spec
	_, err = store.SavePersona(ctx, spec, storage.SaveOptions{
		ID:     req.AgentID,
		Format: storage.FormatJSON,
		Force:  true,
	})
	if err != nil {
		sendError(sc, "Failed to save modified persona: "+err.Error())
		return
	}

	log.Printf("Persona modified for %s: %s", req.AgentID, description)
	sc.WriteJSON(WSMessage{Type: "persona_modified", Payload: map[string]interface{}{
		"agent_id":    req.AgentID,
		"modified":    true,
		"description": description,
	}})
}

// toneDownEyeColor replaces vivid/striking eye color descriptors with softer terms
// It modifies the Visual.BiologicalIdentity.Eyes.Color or Eyes.DistinctiveFeatures
func toneDownEyeColor(spec *persona.PersonaSpec, intensity string) bool {
	// Check both the visual spec and distinctive features
	eyeColor := spec.Visual.BiologicalIdentity.Eyes.Color
	eyeFeatures := spec.Visual.BiologicalIdentity.Eyes.DistinctiveFeatures

	if eyeColor == "" && eyeFeatures == "" {
		return false
	}

	// Define replacements based on intensity
	var replacements map[string]string
	if intensity == "subtle" {
		replacements = map[string]string{
			"striking":      "soft",
			"vivid":         "gentle",
			"bright":        "warm",
			"intense":       "natural",
			"piercing":      "calm",
			"electric":      "muted",
			"brilliant":     "subdued",
			"dazzling":      "understated",
			"mesmerizing":   "pleasant",
			"hypnotic":      "relaxed",
			"captivating":   "appealing",
			"enchanting":    "nice",
			"extraordinary": "ordinary",
		}
	} else { // natural (default)
		replacements = map[string]string{
			"striking":      "natural",
			"vivid":         "natural",
			"bright":        "clear",
			"intense":       "gentle",
			"piercing":      "soft",
			"electric":      "calm",
			"brilliant":     "clear",
			"dazzling":      "pleasant",
			"mesmerizing":   "appealing",
			"hypnotic":      "natural",
			"captivating":   "nice",
			"enchanting":    "pleasant",
			"extraordinary": "normal",
		}
	}

	modified := false

	// Apply to eye color description
	if eyeColor != "" {
		newColor := applyReplacements(eyeColor, replacements)
		if newColor != eyeColor {
			spec.Visual.BiologicalIdentity.Eyes.Color = newColor
			modified = true
		}
	}

	// Apply to distinctive features
	if eyeFeatures != "" {
		newFeatures := applyReplacements(eyeFeatures, replacements)
		if newFeatures != eyeFeatures {
			spec.Visual.BiologicalIdentity.Eyes.DistinctiveFeatures = newFeatures
			modified = true
		}
	}

	return modified
}

// applyReplacements applies case-insensitive replacements while preserving original casing
func applyReplacements(original string, replacements map[string]string) string {
	result := original
	for old, new := range replacements {
		// Case-insensitive replacement
		lowerResult := strings.ToLower(result)
		lowerOld := strings.ToLower(old)
		if idx := strings.Index(lowerResult, lowerOld); idx != -1 {
			// Preserve casing by checking if original was capitalized
			if idx == 0 && len(result) > 0 && result[0] >= 'A' && result[0] <= 'Z' {
				new = strings.ToUpper(string(new[0])) + new[1:]
			}
			result = result[:idx] + new + result[idx+len(old):]
		}
	}
	return result
}

// alignHairAndBrows ensures DNA brow color logically matches hair color
// Note: Since DNA uses categorical types, we log the mismatch but can't directly
// change enum values. For visual spec, we can update string descriptions.
func alignHairAndBrows(spec *persona.PersonaSpec) bool {
	// Get hair color from either DNA or Visual spec
	var hairColor string
	if spec.DNA != nil && spec.DNA.Face.HairColor != "" {
		hairColor = string(spec.DNA.Face.HairColor)
	} else {
		hairColor = spec.Visual.BiologicalIdentity.Hair.CurrentColor
		if hairColor == "" {
			hairColor = spec.Visual.BiologicalIdentity.Hair.NaturalColor
		}
	}

	if hairColor == "" {
		return false
	}

	// Map hair colors to expected brow darkness
	// This helps identify mismatches (e.g., black hair with blonde brows)
	browColorMap := map[string]string{
		"black":        "dark",
		"dark_brown":   "dark brown",
		"dark brown":   "dark brown",
		"brown":        "brown",
		"light_brown":  "light brown",
		"light brown":  "light brown",
		"auburn":       "auburn",
		"red":          "auburn-tinted",
		"ginger":       "ginger-tinted",
		"blonde":       "light brown",
		"dirty blonde": "light brown",
		"platinum":     "light",
		"gray":         "gray",
		"grey":         "grey",
		"white":        "light gray",
		"silver":       "gray",
	}

	// Find matching brow color
	hairColorLower := strings.ToLower(hairColor)
	var expectedBrowColor string
	for key, val := range browColorMap {
		if strings.Contains(hairColorLower, key) {
			expectedBrowColor = val
			break
		}
	}

	if expectedBrowColor == "" {
		// Default: derive from hair color
		expectedBrowColor = hairColorLower
	}

	// Store the expected brow color in a custom field or style signature note
	// Since DNA brows are categorical (shape/thickness), we add a note to the style signature
	signatureMakeup := spec.Visual.StyleSignature.SignatureMakeup
	browNote := fmt.Sprintf("brows match %s hair", expectedBrowColor)

	if signatureMakeup == "" {
		spec.Visual.StyleSignature.SignatureMakeup = browNote
		return true
	}

	// Check if already has brow note
	if strings.Contains(strings.ToLower(signatureMakeup), "brow") {
		return false
	}

	// Append brow note
	spec.Visual.StyleSignature.SignatureMakeup = signatureMakeup + ", " + browNote
	return true
}

func handleRegenerateBanner(sc *SafeConn, req RegenerateBannerRequest) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute)
	defer cancel()

	// Load agent
	agent, err := loadAgent(req.AgentID)
	if err != nil {
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	// Create service
	cfg := personagen.ServiceConfig{
		LaozhangAPIKey: apiKey,
	}
	service, err := personagen.NewService(ctx, cfg, &personagen.NoopProgressReporter{})
	if err != nil {
		sendError(sc, "Failed to create service: "+err.Error())
		return
	}

	// Load persona spec and anchor
	optimizer := storage.NewOptimizer(85, 80)
	store := storage.NewFilesystemStorage(mediaDir, optimizer)
	spec, err := store.LoadPersona(ctx, agent.ID)
	if err != nil {
		sendError(sc, "Failed to load persona: "+err.Error())
		return
	}

	// Load anchor for consistency
	anchorBytes, err := store.LoadAnchor(ctx, agent.ID)
	if err != nil {
		sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]string{
			"message": "Anchor not found - regenerate anchor first",
			"code":    "ANCHOR_MISSING",
		}})
		return
	}
	service.SetAnchor(anchorBytes)

	// Notify that regeneration is starting
	sc.WriteJSON(WSMessage{Type: "banner_regenerating", Payload: map[string]interface{}{
		"agent_id": req.AgentID,
		"style":    req.Style,
	}})

	// Apply style hint if provided
	if req.Style != "" {
		applyBannerStyleHint(spec, req.Style)
	}

	// Delete existing banner
	bannerDir := filepath.Join(mediaDir, agent.ID)
	entries, _ := os.ReadDir(bannerDir)
	for _, entry := range entries {
		if strings.HasPrefix(entry.Name(), "banner") {
			os.Remove(filepath.Join(bannerDir, entry.Name()))
		}
	}

	// Generate banner using the job function
	generateBannerJob(ctx, service, store, spec, agent)
}

// applyBannerStyleHint modifies the persona spec to hint at a banner style
func applyBannerStyleHint(spec *persona.PersonaSpec, style string) {
	// We store the style hint in a way the banner generation can use
	// This could be done via ImageMatrix or a custom field
	// For now, we'll use the scene hint approach by modifying lifestyle context
	switch style {
	case "lifestyle":
		// Emphasize lifestyle/activity context
		if len(spec.Lifestyle.Interests.Passionate) > 0 {
			// Already has lifestyle context, no change needed
		}
	case "portrait":
		// For portrait, we want a simpler background
		// This would need banner-specific handling in GenerateBanner
	case "scenic":
		// Emphasize outdoor/scenic elements
	case "artistic":
		// Allow more creative/artistic interpretation
	}
	// The actual implementation depends on how GenerateBanner constructs prompts
	// This is a placeholder for the style hint mechanism
}

// handleBulkRegenerateVideos regenerates all videos for multiple agents.
// This is useful after fixing accent issues to update all affected agents.
func handleBulkRegenerateVideos(sc *SafeConn, req BulkRegenerateVideosRequest) {
	if len(req.AgentIDs) == 0 {
		sendError(sc, "No agent IDs provided")
		return
	}

	// Notify start of bulk operation
	sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_started", Payload: map[string]interface{}{
		"agent_ids": req.AgentIDs,
		"total":     len(req.AgentIDs),
	}})

	log.Printf("Starting bulk video regeneration for %d agents", len(req.AgentIDs))

	// Process each agent sequentially (to avoid overwhelming the video API)
	for i, agentID := range req.AgentIDs {
		agent, err := loadAgent(agentID)
		if err != nil {
			log.Printf("Agent not found: %s", agentID)
			sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
				"current":    i + 1,
				"total":      len(req.AgentIDs),
				"agent_id":   agentID,
				"agent_name": "",
				"status":     "error",
				"error":      "Agent not found",
			}})
			continue
		}

		// Delete existing videos
		for _, videoURL := range agent.Videos {
			relativePath := strings.TrimPrefix(strings.Split(videoURL, "?")[0], "/media/")
			fullPath := filepath.Join(mediaDir, relativePath)
			os.Remove(fullPath)
		}

		// Clear videos from agent
		updatedAgent, _ := updateAgentField(agentID, func(a *Agent) {
			a.Videos = []string{}
		})
		hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})

		// Regenerate video for this agent
		ctx, cancel := context.WithTimeout(context.Background(), 15*time.Minute)

		cfg := personagen.ServiceConfig{
			LaozhangAPIKey: apiKey,
		}
		service, err := personagen.NewService(ctx, cfg, &personagen.NoopProgressReporter{})
		if err != nil {
			cancel()
			log.Printf("Failed to create service for %s: %v", agent.Name, err)
			sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
				"current":    i + 1,
				"total":      len(req.AgentIDs),
				"agent_id":   agentID,
				"agent_name": agent.Name,
				"status":     "error",
				"error":      err.Error(),
			}})
			continue
		}

		optimizer := storage.NewOptimizer(85, 80)
		store := storage.NewFilesystemStorage(mediaDir, optimizer)
		spec, err := store.LoadPersona(ctx, agentID)
		if err != nil {
			cancel()
			log.Printf("Failed to load persona for %s: %v", agent.Name, err)
			sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
				"current":    i + 1,
				"total":      len(req.AgentIDs),
				"agent_id":   agentID,
				"agent_name": agent.Name,
				"status":     "error",
				"error":      err.Error(),
			}})
			continue
		}

		// Load anchor
		anchorBytes, err := store.LoadAnchor(ctx, agentID)
		if err != nil {
			cancel()
			log.Printf("Anchor not found for %s: %v", agent.Name, err)
			sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
				"current":    i + 1,
				"total":      len(req.AgentIDs),
				"agent_id":   agentID,
				"agent_name": agent.Name,
				"status":     "error",
				"error":      "Anchor not found",
			}})
			continue
		}
		service.SetAnchor(anchorBytes)

		// Notify progress
		sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
			"current":    i + 1,
			"total":      len(req.AgentIDs),
			"agent_id":   agentID,
			"agent_name": agent.Name,
			"status":     "generating",
		}})

		// Generate video
		generateVideoJob(ctx, service, store, spec, agent)
		cancel()

		// Notify completion for this agent
		sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_progress", Payload: map[string]interface{}{
			"current":    i + 1,
			"total":      len(req.AgentIDs),
			"agent_id":   agentID,
			"agent_name": agent.Name,
			"status":     "complete",
		}})

		log.Printf("Video regenerated for %s (%d/%d)", agent.Name, i+1, len(req.AgentIDs))
	}

	// Notify completion of bulk operation
	sc.WriteJSON(WSMessage{Type: "bulk_video_regeneration_complete", Payload: map[string]interface{}{
		"total": len(req.AgentIDs),
	}})
	log.Printf("Bulk video regeneration complete for %d agents", len(req.AgentIDs))
}

func handleBulkDelete(sc *SafeConn, req BulkDeleteRequest) {
	// Build sets of URLs to delete (strip query params for comparison)
	imageURLsToDelete := make(map[string]bool)
	for _, url := range req.ImageURLs {
		imageURLsToDelete[url] = true
	}

	videoURLsToDelete := make(map[string]bool)
	for _, url := range req.VideoURLs {
		urlWithoutQuery := strings.Split(url, "?")[0]
		videoURLsToDelete[urlWithoutQuery] = true
	}

	var imagesDeleted, videosDeleted int

	// Update agent atomically and delete files
	updatedAgent, err := updateAgentField(req.AgentID, func(a *Agent) {
		// Delete image files and filter array
		newImages := []string{}
		for _, imgURL := range a.Images {
			if imageURLsToDelete[imgURL] {
				// Delete the actual files (webp, avif, caption)
				relativePath := strings.TrimPrefix(imgURL, "/media/")
				basePath := filepath.Join(mediaDir, relativePath)
				baseWithoutExt := strings.TrimSuffix(basePath, filepath.Ext(basePath))

				os.Remove(basePath)                    // .webp
				os.Remove(baseWithoutExt + ".avif")    // .avif
				os.Remove(baseWithoutExt + ".caption") // .caption

				imagesDeleted++
				continue
			}
			newImages = append(newImages, imgURL)
		}
		a.Images = newImages

		// Delete video files and filter array
		newVideos := []string{}
		for _, vidURL := range a.Videos {
			vidWithoutQuery := strings.Split(vidURL, "?")[0]
			if videoURLsToDelete[vidWithoutQuery] {
				// Delete the actual video file
				relativePath := strings.TrimPrefix(vidWithoutQuery, "/media/")
				fullPath := filepath.Join(mediaDir, relativePath)
				os.Remove(fullPath)

				videosDeleted++
				continue
			}
			newVideos = append(newVideos, vidURL)
		}
		a.Videos = newVideos
	})

	if err != nil {
		sendError(sc, "Agent not found: "+req.AgentID)
		return
	}

	log.Printf("Bulk deleted %d images and %d videos from %s", imagesDeleted, videosDeleted, updatedAgent.Name)
	hub.Broadcast(WSMessage{Type: "agent_updated", Payload: updatedAgent})
	sc.WriteJSON(WSMessage{Type: "bulk_deleted", Payload: map[string]int{
		"images_deleted": imagesDeleted,
		"videos_deleted": videosDeleted,
	}})
}

func createJob(agent *Agent, jobType string, total int) *MediaJob {
	job := &MediaJob{
		ID:        uuid.New().String(),
		AgentID:   agent.ID,
		AgentName: agent.Name,
		Type:      jobType,
		Status:    "pending",
		Progress:  0,
		Total:     total,
	}
	jobsMu.Lock()
	jobs[job.ID] = job
	jobsMu.Unlock()
	hub.Broadcast(WSMessage{Type: "job_update", Payload: job})
	return job
}

func updateJob(job *MediaJob) {
	jobsMu.Lock()
	jobs[job.ID] = job
	jobsMu.Unlock()
	hub.Broadcast(WSMessage{Type: "job_update", Payload: job})

	// Auto-remove completed/error jobs after 30 seconds
	if job.Status == "complete" || job.Status == "error" {
		go func(jobID string) {
			time.Sleep(30 * time.Second)
			jobsMu.Lock()
			delete(jobs, jobID)
			jobsMu.Unlock()
			hub.Broadcast(WSMessage{Type: "job_removed", Payload: map[string]string{"id": jobID}})
		}(job.ID)
	}
}

func clearCompletedJobs() {
	jobsMu.Lock()
	toRemove := []string{}
	for id, job := range jobs {
		if job.Status == "complete" || job.Status == "error" {
			toRemove = append(toRemove, id)
		}
	}
	for _, id := range toRemove {
		delete(jobs, id)
	}
	jobsMu.Unlock()

	for _, id := range toRemove {
		hub.Broadcast(WSMessage{Type: "job_removed", Payload: map[string]string{"id": id}})
	}
}

func extractTags(spec *persona.PersonaSpec) []string {
	tags := []string{}

	// Add archetype as a tag
	if spec.Psychology.DatingPersonality.Archetype != "" {
		tags = append(tags, strings.ToLower(spec.Psychology.DatingPersonality.Archetype))
	}

	// Add passionate interests as tags
	for _, interest := range spec.Lifestyle.Interests.Passionate {
		if len(tags) < 8 {
			tags = append(tags, strings.ToLower(interest))
		}
	}

	// Fill remaining slots with casual interests
	for _, interest := range spec.Lifestyle.Interests.Casual {
		if len(tags) < 8 {
			tags = append(tags, strings.ToLower(interest))
		}
	}

	return tags
}

func saveAgentJSON(agent *Agent) error {
	data, _ := json.MarshalIndent(agent, "", "  ")
	return os.WriteFile(filepath.Join(agentsDir, agent.ID+".json"), data, 0644)
}

// updateAgentField atomically reads the agent, applies the update function, and saves.
// This prevents race conditions when multiple jobs update the same agent concurrently.
func updateAgentField(agentID string, updateFn func(*Agent)) (*Agent, error) {
	agentsMu.Lock()
	defer agentsMu.Unlock()

	agent, err := loadAgent(agentID)
	if err != nil {
		return nil, err
	}
	updateFn(agent)
	if err := saveAgentJSON(agent); err != nil {
		return nil, err
	}
	return agent, nil
}

func loadAgent(id string) (*Agent, error) {
	data, err := os.ReadFile(filepath.Join(agentsDir, id+".json"))
	if err != nil {
		return nil, err
	}
	var agent Agent
	if err := json.Unmarshal(data, &agent); err != nil {
		return nil, err
	}
	return &agent, nil
}

func listAgents(w http.ResponseWriter, r *http.Request) {
	// Parse pagination params
	limit := 20
	offset := 0
	if l := r.URL.Query().Get("limit"); l != "" {
		if parsed, err := strconv.Atoi(l); err == nil && parsed > 0 {
			limit = parsed
		}
	}
	if o := r.URL.Query().Get("offset"); o != "" {
		if parsed, err := strconv.Atoi(o); err == nil && parsed >= 0 {
			offset = parsed
		}
	}

	// Load all agents
	files, _ := os.ReadDir(agentsDir)
	allAgents := []Agent{}
	for _, f := range files {
		if !strings.HasSuffix(f.Name(), ".json") {
			continue
		}
		data, err := os.ReadFile(filepath.Join(agentsDir, f.Name()))
		if err != nil {
			continue
		}
		var agent Agent
		if json.Unmarshal(data, &agent) == nil {
			allAgents = append(allAgents, agent)
		}
	}

	// Sort by created_at descending (newest first)
	sort.Slice(allAgents, func(i, j int) bool {
		return allAgents[i].CreatedAt.After(allAgents[j].CreatedAt)
	})

	// Apply pagination
	total := len(allAgents)
	end := offset + limit
	if end > total {
		end = total
	}
	start := offset
	if start > total {
		start = total
	}
	paged := allAgents[start:end]

	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(map[string]interface{}{
		"agents":   paged,
		"total":    total,
		"has_more": end < total,
	})
}

func getAgent(w http.ResponseWriter, r *http.Request) {
	id := strings.TrimPrefix(r.URL.Path, "/agents/")
	data, err := os.ReadFile(filepath.Join(agentsDir, id+".json"))
	if err != nil {
		http.Error(w, "Not found", http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write(data)
}

func sendError(sc *SafeConn, msg string) {
	sc.WriteJSON(WSMessage{Type: "error", Payload: map[string]string{"message": msg}})
}

// wsProgressReporter implements personagen.ProgressReporter for WebSocket updates
type wsProgressReporter struct {
	send func(WSMessage)
}

func (r *wsProgressReporter) OnProgress(stage, message string) {
	r.send(WSMessage{Type: "progress", Payload: map[string]string{"stage": stage, "message": message}})
}