package main

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

	"frame"
	"gocv.io/x/gocv"
	"maunium.net/go/mautrix/id"
)

func main() {
	var (
		listenAddr      string
		passphrase      string
		authToken       string
		matrixHS        string
		matrixUser      string
		matrixTokenFile string
		clipsRoom       string
		detectionsRoom  string
		viewRoom        string
		pbkdf2Iter      int
	)

	flag.StringVar(&listenAddr, "listen", ":8082", "address to listen for clips from terp")
	flag.StringVar(&passphrase, "passphrase", "", "decryption passphrase (empty = no decryption)")
	flag.StringVar(&authToken, "auth-token", "", "shared pipeline auth token (or OKO_AUTH_TOKEN)")
	flag.StringVar(&matrixHS, "matrix-homeserver", "", "Matrix homeserver URL (required)")
	flag.StringVar(&matrixUser, "matrix-user", "", "Matrix user ID, e.g. @oko:example.org (required)")
	flag.StringVar(&matrixTokenFile, "matrix-token-file", "", "file containing the Matrix access token (or OKO_MATRIX_TOKEN env)")
	flag.StringVar(&clipsRoom, "clips-room", "", "room ID for all motion clips (required)")
	flag.StringVar(&detectionsRoom, "detections-room", "", "room ID for clips with detections (required)")
	flag.StringVar(&viewRoom, "view-room", "", "room ID for clips with minimal caption only, no JSON metadata (required)")
	flag.IntVar(&pbkdf2Iter, "pbkdf2-iters", frame.DefaultPBKDF2Iter, "PBKDF2 iterations for clip encryption (1..1000000)")
	flag.Parse()

	if pbkdf2Iter < 1 || pbkdf2Iter > frame.MaxPBKDF2Iter {
		log.Fatalf("pbkdf2-iters must be in [1, %d], got %d", frame.MaxPBKDF2Iter, pbkdf2Iter)
	}
	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)

	token := os.Getenv("OKO_MATRIX_TOKEN")
	if matrixTokenFile != "" {
		data, err := os.ReadFile(matrixTokenFile)
		if err != nil {
			log.Fatalf("Read matrix token file: %v", err)
		}
		token = strings.TrimSpace(string(data))
	}
	if matrixHS == "" || matrixUser == "" || token == "" {
		log.Fatal("matrix-homeserver, matrix-user, and a token (OKO_MATRIX_TOKEN or -matrix-token-file) are required")
	}

	ms, err := newMatrixStore(matrixHS, matrixUser, token, clipsRoom, detectionsRoom, viewRoom)
	if err != nil {
		log.Fatalf("Initialize Matrix backend: %v", err)
	}
	log.Printf("Matrix backend ready")

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

	ln, err := net.Listen("tcp", listenAddr)
	if err != nil {
		log.Fatalf("Listen %s: %v", listenAddr, err)
	}
	log.Printf("Listening on %s", listenAddr)

	var wg sync.WaitGroup
	rejected := frame.NewMalformedAttempts()
	// Cap concurrent connections so a flood (or slowloris-style connection
	// that parks on the auth handshake) cannot exhaust goroutines/FDs.
	connSem := make(chan struct{}, 16)
	// Cap concurrent Matrix ingest so a flood of clips cannot spawn unbounded
	// encode+upload work. Deliberately acquired only inside the authenticated
	// callback, so an unauthenticated connection cannot occupy an ingest slot.
	sem := make(chan struct{}, 8)
	go func() {
		<-ctx.Done()
		ln.Close()
	}()

	for {
		conn, err := ln.Accept()
		if err != nil {
			if ctx.Err() != nil {
				break
			}
			log.Printf("Accept error: %v", err)
			continue
		}
		select {
		case connSem <- struct{}{}:
		default:
			log.Printf("Rejecting connection from %s: too many concurrent connections", conn.RemoteAddr())
			conn.Close()
			continue
		}
		wg.Add(1)
		go func(c net.Conn) {
			defer func() { <-connSem }()
			defer wg.Done()
			frame.HandleClipConn(c, rejected, func(clp frame.Clip) {
				sem <- struct{}{}
				defer func() { <-sem }()
				handleClip(ctx, ms, passphrase, clp)
			})
		}(conn)
	}

	if !frame.WaitGroupTimeout(&wg, 5*time.Second) {
		log.Printf("In-flight handlers still working; exiting anyway")
	}
	if n := rejected.Total(); n > 0 {
		log.Printf("Rejected %d malformed clips from the network", n)
	}
	log.Println("Exiting")
}

func handleClip(parent context.Context, ms *matrixStore, passphrase string, clp frame.Clip) {
	targets := routeTargets(&clp, ms.clipsRoom, ms.detectionsRoom)
	if len(targets) == 0 {
		log.Printf("Dropped clip with no motion (source %s)", clp.SourceData)
		return
	}

	if passphrase != "" {
		if err := clp.Decrypt(passphrase); err != nil {
			log.Printf("Decrypt error: %v", err)
			return
		}
	}

	frames := clp.Sublimate()
	if len(frames) == 0 {
		return
	}

	// Timebox encode + upload + posts together; this is the only archival hop,
	// so a transient homeserver blip should be retried rather than dropping
	// the footage.
	ctx, cancel := context.WithTimeout(parent, 5*time.Minute)
	defer cancel()

	fps := computeFPS(clp.Timestamps)
	mp4, err := encodeMP4(frames, int(clp.Width), int(clp.Height), fps)
	if err != nil {
		log.Printf("Encode clip error: %v", err)
		return
	}

	filename := fmt.Sprintf("%x-%d.mp4", clp.Guids[0], clp.Timestamps[0])
	var mxc id.ContentURI
	if err := withRetry(ctx, 3, func() error {
		mxc, err = ms.uploadClip(ctx, filename, mp4)
		return err
	}); err != nil {
		log.Printf("Upload clip media: %v", err)
		return
	}

	failed := 0
	for _, roomID := range targets {
		var eventID id.EventID
		if err := withRetry(ctx, 2, func() error {
			eventID, err = ms.sendClipToRoom(ctx, roomID, clp, filename, mxc, len(mp4), fps)
			return err
		}); err != nil {
			failed++
			log.Printf("Post clip to %s: %v", roomID, err)
			continue
		}
		log.Printf("Stored clip (%d frames, motion=%d%%, %.1f fps) in %s as %s",
			len(frames), clp.GetHighestMotion(), fps, roomID, eventID)
	}

	var viewID id.EventID
	if err := withRetry(ctx, 2, func() error {
		viewID, err = ms.sendClipView(ctx, ms.viewRoom, clp, mxc, len(mp4), fps)
		return err
	}); err != nil {
		failed++
		log.Printf("Post clip to view room: %v", err)
	} else {
		log.Printf("Stored clip (view, minimal) in %s as %s", ms.viewRoom, viewID)
	}

	if failed > 0 {
		log.Printf("Failed to post clip to %d of %d target rooms", failed, len(targets)+1)
	}
}

// withRetry runs fn up to attempts times with capped exponential backoff,
// returning early when ctx is done.
func withRetry(ctx context.Context, attempts int, fn func() error) error {
	if attempts < 1 {
		attempts = 1
	}
	var lastErr error
	for i := 1; i <= attempts; i++ {
		if err := fn(); err != nil {
			lastErr = err
			delay := 500 * time.Millisecond * time.Duration(1<<(i-1))
			if delay > 5*time.Second {
				delay = 5 * time.Second
			}
			select {
			case <-ctx.Done():
				return lastErr
			case <-time.After(delay):
			}
			continue
		}
		return nil
	}
	return lastErr
}

func computeFPS(timestamps []uint64) float64 {
	if len(timestamps) < 2 {
		return 30.0
	}

	// Average of the consecutive inter-frame deltas, so a single irregular
	// pair of timestamps cannot skew the reported rate.
	var sum uint64
	deltas := 0
	for i := 1; i < len(timestamps); i++ {
		if timestamps[i] <= timestamps[i-1] {
			continue
		}
		sum += timestamps[i] - timestamps[i-1]
		deltas++
	}
	if deltas == 0 {
		return 30.0
	}

	avg := float64(sum) / float64(deltas)
	fps := 1e9 / avg
	if fps <= 0 || fps > 120 {
		return 30.0
	}
	return math.Round(fps*100) / 100
}

// encodeMP4 renders the clip's frames to an H.264 MP4 in a temporary file and
// returns the bytes. The temp file is removed before returning.
func encodeMP4(frames []frame.Frame, width, height int, fps float64) ([]byte, error) {
	tmp, err := os.CreateTemp("", "oko-clip-*.mp4")
	if err != nil {
		return nil, err
	}
	path := tmp.Name()
	tmp.Close()
	defer os.Remove(path)

	if err := writeVideo(path, frames, width, height, fps); err != nil {
		return nil, err
	}
	return os.ReadFile(path)
}

func writeVideo(path string, frames []frame.Frame, width, height int, fps float64) error {
	isColor := true
	if len(frames) > 0 && frames[0].Channels == 1 {
		isColor = false
	}

	var vw *gocv.VideoWriter
	var err error

	for _, codec := range []string{"avc1", "H264", "MJPG"} {
		vw, err = gocv.VideoWriterFile(path, codec, fps, width, height, isColor)
		if err == nil && vw.IsOpened() {
			break
		}
		if vw != nil {
			vw.Close()
		}
		vw = nil
	}

	if vw == nil {
		return fmt.Errorf("failed to open video writer with any codec: %v", err)
	}
	defer vw.Close()

	written := 0
	for i := range frames {
		mat, err := frames[i].ToMat()
		if err != nil {
			log.Printf("ToMat error frame %d: %v", i, err)
			continue
		}
		if err := vw.Write(mat); err != nil {
			mat.Close()
			return fmt.Errorf("write frame %d: %w", i, err)
		}
		mat.Close()
		written++
	}

	if written == 0 {
		return fmt.Errorf("no frames could be encoded")
	}
	if written < len(frames) {
		log.Printf("Encoded %d of %d frames (partial clip)", written, len(frames))
	}

	return nil
}