package fall

import (
	"fmt"
	"math"
	"sort"

	"silverpose/v2/internal/pose"
)

type track struct {
	centerX    float32
	centerY    float32
	lastSeenAt float64
}

type tracker struct {
	maxMatchDistanceRatio float64
	maxAgeSeconds         float64
	tracks                map[string]track
	nextTrackNumber       int
}

func newTracker(maxMatchDistanceRatio, maxAgeSeconds float64) (*tracker, error) {
	if maxMatchDistanceRatio <= 0 || maxMatchDistanceRatio > 1 {
		return nil, fmt.Errorf("max match distance ratio must be in (0, 1]")
	}
	if maxAgeSeconds <= 0 {
		return nil, fmt.Errorf("max age seconds must be positive")
	}
	return &tracker{
		maxMatchDistanceRatio: maxMatchDistanceRatio,
		maxAgeSeconds:         maxAgeSeconds,
		tracks:                make(map[string]track),
		nextTrackNumber:       1,
	}, nil
}

func (tracker *tracker) update(poses []pose.PersonPose, now float64, width, height int) ([]TrackedPose, error) {
	if width <= 0 || height <= 0 {
		return nil, fmt.Errorf("frame size must contain positive width and height")
	}
	tracker.expire(now)
	availableIDs := make(map[string]bool, len(tracker.tracks))
	for trackID := range tracker.tracks {
		availableIDs[trackID] = true
	}
	tracked := make([]TrackedPose, 0, len(poses))
	for _, person := range poses {
		centerX := (person.Box.Left + person.Box.Right) / 2
		centerY := (person.Box.Top + person.Box.Bottom) / 2
		trackID := tracker.nearestAvailable(centerX, centerY, availableIDs, width, height)
		if trackID == "" {
			trackID = fmt.Sprintf("P-%04d", tracker.nextTrackNumber)
			tracker.nextTrackNumber++
		} else {
			delete(availableIDs, trackID)
		}
		tracker.tracks[trackID] = track{centerX: centerX, centerY: centerY, lastSeenAt: now}
		tracked = append(tracked, TrackedPose{TrackID: trackID, Pose: person})
	}
	return tracked, nil
}

func (tracker *tracker) nearestAvailable(centerX, centerY float32, availableIDs map[string]bool, width, height int) string {
	ids := make([]string, 0, len(availableIDs))
	for trackID := range availableIDs {
		ids = append(ids, trackID)
	}
	sort.Strings(ids)
	closestID := ""
	closestDistance := math.Inf(1)
	for _, trackID := range ids {
		previous := tracker.tracks[trackID]
		distance := math.Hypot(
			float64(centerX-previous.centerX)/float64(width),
			float64(centerY-previous.centerY)/float64(height),
		)
		if distance <= tracker.maxMatchDistanceRatio && distance < closestDistance {
			closestID = trackID
			closestDistance = distance
		}
	}
	return closestID
}

func (tracker *tracker) expire(now float64) {
	for trackID, current := range tracker.tracks {
		if now-current.lastSeenAt > tracker.maxAgeSeconds {
			delete(tracker.tracks, trackID)
		}
	}
}