package main

import (
	"context"
	"flag"
	"fmt"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
	"time"

	"frame"
)

// camHandle captures the lifecycle state of one running camera so the
// rebalance loop and the camera's own goroutine can safely agree on removal
// (comparing handles by identity avoids deleting a replacement camera).
type camHandle struct {
	ctx    context.Context
	cancel context.CancelFunc
}

func main() {
	var (
		mofinAddr      string
		camMin         int
		camMax         int
		clipDur        time.Duration
		rescanInterval time.Duration
		username       string
		cameraID       string
		passphrase     string
		authToken      string
		videoFPS       int
		videoWidth     int
		videoHeight    int
		pbkdf2Iter     int
	)

	flag.StringVar(&mofinAddr, "mofin", "localhost:8083", "mofin motion detection address")
	flag.IntVar(&camMin, "min-cam", 0, "minimum camera index to scan")
	flag.IntVar(&camMax, "max-cam", 10, "maximum camera index to scan")
	flag.DurationVar(&clipDur, "clip-duration", 10*time.Second, "duration of each clip")
	flag.DurationVar(&rescanInterval, "rescan-interval", 10*time.Second, "interval between hotplug rescans")
	flag.StringVar(&username, "user", "oko", "username for frame source data")
	flag.StringVar(&cameraID, "camera-id", "", "camera identifier prefix (required)")
	flag.StringVar(&passphrase, "passphrase", "", "encryption passphrase (empty = no encryption)")
	flag.StringVar(&authToken, "auth-token", "", "shared pipeline auth token (or OKO_AUTH_TOKEN)")
	flag.IntVar(&videoFPS, "video-fps", 5, "max frames per second read from the camera and stored in clips (0 = device native rate)")
	flag.IntVar(&videoWidth, "video-width", 0, "capture width in pixels (0 = device default)")
	flag.IntVar(&videoHeight, "video-height", 0, "capture height in pixels (0 = device default)")
	flag.IntVar(&pbkdf2Iter, "pbkdf2-iters", frame.DefaultPBKDF2Iter, "PBKDF2 iterations for clip encryption (1..1000000)")
	flag.Parse()

	if cameraID == "" {
		fmt.Fprintln(os.Stderr, "-camera-id is required")
		flag.Usage()
		os.Exit(1)
	}

	if pbkdf2Iter < 1 || pbkdf2Iter > frame.MaxPBKDF2Iter {
		log.Fatalf("pbkdf2-iters must be in [1, %d], got %d", frame.MaxPBKDF2Iter, pbkdf2Iter)
	}
	if videoFPS < 0 || videoFPS > 60 {
		log.Fatalf("video-fps must be in [0, 60], got %d", videoFPS)
	}
	if videoWidth < 0 || videoHeight < 0 {
		log.Fatalf("video dimensions must be >= 0, got %dx%d", videoWidth, videoHeight)
	}
	frame.SetDefaultPBKDF2Iter(pbkdf2Iter)

	tok := authToken
	if tok == "" {
		tok = os.Getenv("OKO_AUTH_TOKEN")
	}
	if tok == "" {
		log.Fatal("authentication required: set -auth-token or OKO_AUTH_TOKEN")
	}
	frame.SetAuthToken(tok)

	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer cancel()

	var (
		cameras   = make(map[int]*camHandle)
		camerasMu sync.Mutex
		wg        sync.WaitGroup
	)

	// Hotplug loop: scan + rebalance continuously. It lives behind its own
	// WaitGroup so shutdown can wait for it to stop scheduling new cameras
	// BEFORE wg.Wait() (rebalance is the only caller of wg.Add; waiting on it
	// first removes the Add/Wait misuse race).
	var rebalanceWg sync.WaitGroup
	rebalanceWg.Add(1)
	go func() {
		defer rebalanceWg.Done()
		ticker := time.NewTicker(rescanInterval)
		defer ticker.Stop()

		for {
			rebalance(ctx, &wg, &camerasMu, cameras, camMin, camMax, mofinAddr, clipDur, username, cameraID, passphrase, videoFPS, videoWidth, videoHeight)

			select {
			case <-ctx.Done():
				return
			case <-ticker.C:
			}
		}
	}()

	<-ctx.Done()
	log.Println("Shutting down...")

	camerasMu.Lock()
	for _, h := range cameras {
		h.cancel()
	}
	camerasMu.Unlock()

	rebalanceWg.Wait()

	done := make(chan struct{})
	go func() {
		wg.Wait()
		close(done)
	}()
	select {
	case <-done:
		log.Println("All cameras stopped")
	case <-time.After(5 * time.Second): // TODO: Make configurable
		log.Println("Forced shutdown after timeout")
	}
}

func rebalance(ctx context.Context, wg *sync.WaitGroup, mu *sync.Mutex, cameras map[int]*camHandle, camMin, camMax int, mofinAddr string, clipDur time.Duration, username, cameraID, passphrase string, videoFPS, videoWidth, videoHeight int) {
	if ctx.Err() != nil {
		return
	}

	mu.Lock()
	inUse := make(map[int]struct{}, len(cameras))
	for idx := range cameras {
		inUse[idx] = struct{}{}
	}
	mu.Unlock()

	// Probe only devices not already streaming: opening a busy capture handle
	// can transiently report a zero size and flap between teardown and
	// restart. A camera that dies is detected by its own goroutine (read
	// failure) and removed from the map, after which it is probed again.
	indices, err := frame.ScanSkipExcluding(camMin, camMax, 3, inUse)
	if err != nil {
		return
	}

	mu.Lock()
	defer mu.Unlock()

	// Stop tracked cameras that have disappeared. Cameras currently streaming
	// are excluded from the scan, so treat them as still present; vanished
	// ones leave the map themselves via the camera goroutine.
	for idx, h := range cameras {
		if _, busy := inUse[idx]; busy {
			continue
		}
		if !contains(indices, idx) {
			log.Printf("Camera %d: disconnected", idx)
			h.cancel()
			delete(cameras, idx)
		}
	}

	// Start cameras that have appeared
	for _, idx := range indices {
		if _, exists := cameras[idx]; !exists {
			camCtx, c := context.WithCancel(ctx)
			h := &camHandle{ctx: camCtx, cancel: c}
			cameras[idx] = h
			wg.Add(1)
			go func(i int, handle *camHandle) {
				defer wg.Done()
				log.Printf("Camera %d: connected", i)
				runCamera(handle.ctx, i, mofinAddr, clipDur, username, cameraID, passphrase, videoFPS, videoWidth, videoHeight)
				log.Printf("Camera %d: stopped", i)
				mu.Lock()
				if cameras[i] == handle {
					delete(cameras, i)
				}
				mu.Unlock()
			}(idx, h)
		}
	}
}

func contains(slice []int, val int) bool {
	for _, v := range slice {
		if v == val {
			return true
		}
	}
	return false
}

// ThrottleFrames forwards the latest frame received from src to dst at a
// maximum rate of fps frames per second, so a slow reader only ever sees the
// newest frame. If the source closes, the function returns.
func ThrottleFrames(ctx context.Context, src <-chan frame.Frame, dst chan<- frame.Frame, fps int) {
	interval := time.Second / time.Duration(fps)
	if fps <= 0 {
		interval = time.Second / 5
	}

	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	var (
		lastFrame frame.Frame
		haveFrame bool
		frameMu   sync.Mutex
	)

	go func() {
		for {
			select {
			case <-ctx.Done():
				return
			case f, ok := <-src:
				if !ok {
					return
				}
				frameMu.Lock()
				lastFrame = f
				haveFrame = true
				frameMu.Unlock()
			}
		}
	}()

	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			frameMu.Lock()
			if !haveFrame {
				frameMu.Unlock()
				continue
			}
			f := lastFrame
			frameMu.Unlock()

			select {
			case dst <- f:
			default:
			}
		}
	}
}

func runCamera(ctx context.Context, camIndex int, mofinAddr string, clipDur time.Duration, username, cameraID, passphrase string, videoFPS, videoWidth, videoHeight int) {
	frames := make(chan frame.Frame, 60)

	if err := frame.Capture(ctx, camIndex, videoFPS, videoWidth, videoHeight, frames); err != nil {
		log.Printf("Camera %d: %v", camIndex, err)
		return
	}
	log.Printf("Camera %d: started", camIndex)

	sourceData := fmt.Sprintf("%s--%s/%d", username, cameraID, camIndex)

	// Throttle to the configured FPS, then build clips from the sampled frames.
	// Clips are handed to a bounded outbound queue drained by a sender
	// goroutine, so a slow or downed mofin never stalls the capture path —
	// frames keep being sampled and the newest frames are preserved.
	clipFrames := make(chan frame.Frame, 60)
	var throttleWg sync.WaitGroup
	throttleWg.Add(1)
	go func() {
		defer throttleWg.Done()
		ThrottleFrames(ctx, frames, clipFrames, videoFPS)
	}()

	outQueue := make(chan frame.Clip, 4)
	var senderWg sync.WaitGroup
	senderWg.Add(1)
	go func() {
		defer senderWg.Done()
		for clp := range outQueue {
			if err := clp.SendClipRetry(mofinAddr, 6); err != nil {
				log.Printf("Camera %d: send to mofin failed, clip dropped: %v", camIndex, err)
			} else {
				log.Printf("Camera %d: sent clip with %d frames to %s", camIndex, len(clp.PixelMats), mofinAddr)
			}
		}
	}()

	clipWg := sync.WaitGroup{}
	clipWg.Add(1)
	go func() {
		defer clipWg.Done()
		runClipBuilder(ctx, camIndex, clipDur, sourceData, passphrase, clipFrames, outQueue)
	}()

	clipWg.Wait()
	close(outQueue)
	senderWg.Wait()
	throttleWg.Wait()
}

func runClipBuilder(ctx context.Context, camIndex int, clipDur time.Duration, sourceData, passphrase string, frames <-chan frame.Frame, outQueue chan<- frame.Clip) {
	ticker := time.NewTicker(clipDur)
	defer ticker.Stop()

	var buf []frame.Frame
	var pixelBuf [][]byte
	var guidBuf [][]byte
	var tsBuf []uint64

	flush := func() {
		if len(buf) == 0 {
			return
		}

		clip := frame.Clip{
			PixelMats:  pixelBuf,
			Guids:      guidBuf,
			Timestamps: tsBuf,
			Width:      buf[0].Width,
			Height:     buf[0].Height,
			Types:      buf[0].GocvImageType,
			Channels:   buf[0].Channels,
			SourceData: sourceData,
		}

		if passphrase != "" {
			if err := clip.Encrypt(passphrase); err != nil {
				log.Printf("Camera %d: encrypt error: %v", camIndex, err)
			}
		}

		buf = buf[:0]
		pixelBuf = pixelBuf[:0]
		guidBuf = guidBuf[:0]
		tsBuf = tsBuf[:0]

		if len(clip.PixelMats) == 0 {
			return
		}

		// Non-blocking enqueue: the sender goroutine owns retries and network
		// blocking. If the queue is full (mofin down for a while), drop this
		// clip rather than stall capture and lose newest frames.
		select {
		case outQueue <- clip:
		default:
			log.Printf("Camera %d: outbound queue full, dropping clip (%d frames)", camIndex, len(clip.PixelMats))
			// Flush in reverse on drop would hold newer frames back; the
			// queue drains continuously and the next clip is a full window.
		}
	}
	defer flush()

	for {
		select {
		case <-ctx.Done():
			flush()
			return
		case f := <-frames:
			buf = append(buf, f)
			pixelBuf = append(pixelBuf, f.PixelBytes)
			guidBuf = append(guidBuf, f.Guid)
			tsBuf = append(tsBuf, f.Timestamp)
		case <-ticker.C:
			flush()
		}
	}
}