// File: planner/main.go
package main

import (
	"encoding/json"
	"flag"
	"io"
	"math"
	"os"
	"sort"
	"time"
)

type AvailabilitySlot struct {
	DayOfWeek int    `json:"dayOfWeek"` // Go weekday: 0=Sunday..6=Saturday
	Start     string `json:"start"`     // "HH:MM"
	End       string `json:"end"`       // "HH:MM"
}

type TimeOffPeriod struct {
	FromISO string `json:"fromISO"`
	ToISO   string `json:"toISO"`
}

type UserInput struct {
	UserID               int                `json:"user_id"`
	Name                 string             `json:"name"`
	Roles                []int              `json:"roles"`
	RequiredHoursPerWeek float64            `json:"required_hours_per_week"`
	Vacations            []TimeOffPeriod    `json:"vacations"`
	SickDays             []TimeOffPeriod    `json:"sick_days"`
	Availability         []AvailabilitySlot `json:"availability"`
	LastEndISO           *string            `json:"last_end_iso"`
}

type RoleNeed struct {
	RoleID           int      `json:"role_id"`
	Count            int      `json:"count"`
	PreAssignedUsers [](*int) `json:"pre_assigned_users"`
}

type ShiftTemplateInput struct {
	ID                int        `json:"zst_id"`
	Name              string     `json:"zst_name"`
	StartTime         string     `json:"zst_start_time"` // "HH:MM"
	EndTime           string     `json:"zst_end_time"`   // "HH:MM"
	DayMask           int        `json:"zst_day_mask"`
	MinWorkers        int        `json:"zst_min_workers"`
	MaxWorkers        int        `json:"zst_max_workers"`
	MustHaveRoles     []RoleNeed `json:"zst_must_have_roles"`
	NiceToHaveRoles   []RoleNeed `json:"zst_nice_to_have_roles"`
	Active            bool       `json:"zst_active"`
	Priority          int        `json:"zst_priority"`
	HolidayMode       bool       `json:"zst_holiday_mode"`
	ParentTemplateID  *int       `json:"zst_parent_template_id,omitempty"`
	HolidayIds        []int      `json:"holiday_ids"`
}

type RulesInput struct {
	MinRestHours        int  `json:"min_rest_hours"`
	MaxDailyHours       int  `json:"max_daily_hours"`
	MaxHoursPerWeek     int  `json:"max_hours_per_week"`
	MaxConsecutiveDays  int  `json:"max_consecutive_days"`
	EnforceAvailability bool `json:"enforce_availability"`
}

type FairnessInput struct {
	TotalShiftTolerancePct float64 `json:"total_shift_tolerance_pct"`
	MinFreeWeekends        int     `json:"min_free_weekends"`
	WeekendDays            []int   `json:"weekend_days"` // [6,0] => Sat+Sun
	FreeWeekendDefinition  string  `json:"free_weekend_definition"`
}

type HolidayInput struct {
	ID              int    `json:"id"`
	CalculationType string `json:"calculation_type"`
	Month           int    `json:"month"`
	Day             int    `json:"day"`
	EasterOffset    int    `json:"easter_offset"`
}

type ExistingShiftInput struct {
	ShiftID    int    `json:"shift_id"`
	Date       string `json:"date"`       // YYYY-MM-DD
	TemplateID int    `json:"template_id"`
	StartTime  string `json:"start_time"` // HH:MM
	EndTime    string `json:"end_time"`   // HH:MM
}

type PlanPayload struct {
	DepartmentID   int                   `json:"department_id"`
	Year           int                   `json:"year"`
	Month          int                   `json:"month"`
	Users          []UserInput           `json:"users"`
	ShiftTemplates []ShiftTemplateInput  `json:"shift_templates"`
	Holidays       []HolidayInput        `json:"holidays"`
	ExistingShifts []ExistingShiftInput  `json:"existing_shifts"`
	Timezone       string                `json:"timezone"`
	Rules          RulesInput            `json:"rules"`
	Fairness       FairnessInput         `json:"fairness"`
}

type PlanRequest struct {
	Status  string      `json:"status,omitempty"`
	Message string      `json:"message,omitempty"`
	Payload PlanPayload `json:"payload"`
	Note    string      `json:"note,omitempty"`
}

type Window struct {
	StartISO string `json:"startISO"`
	EndISO   string `json:"endISO"`
}

type AssignmentOut struct {
	Date       string  `json:"date"`
	TemplateID int     `json:"template_id"`
	StartISO   string  `json:"startISO"`
	EndISO     string  `json:"endISO"`
	RoleID     int     `json:"role_id"`
	SlotIndex  int     `json:"slot_index"`
	UserID     *int    `json:"user_id,omitempty"`
	Note       *string `json:"note,omitempty"`
}

type UserSummaryOut struct {
	UserID          int                `json:"user_id"`
	Name            string             `json:"name"`
	TotalShifts     int                `json:"total_shifts"`
	TotalMinutes    int                `json:"total_minutes"`
	TemplateCounts  map[string]int     `json:"template_counts"`
	TemplatePercent map[string]float64 `json:"template_percent"`
	FreeWeekends    int                `json:"free_weekends"`
	WeekendWorked   int                `json:"weekend_worked"`
}

type TotalEvenness struct {
	Avg        float64 `json:"avg"`
	Min        int     `json:"min"`
	Max        int     `json:"max"`
	MaxDiffPct float64 `json:"max_diff_pct"`
}

type TemplateEvennessItem struct {
	TemplateID     int     `json:"template_id"`
	Avg            float64 `json:"avg"`
	Min            int     `json:"min"`
	Max            int     `json:"max"`
	MaxDiffFromAvg float64 `json:"max_diff_from_avg"`
}

type FairnessOut struct {
	TolerancePct      float64                `json:"tolerance_pct"`
	OverallScore      float64                `json:"overall_score"`
	UserTotalEvenness TotalEvenness          `json:"user_total_evenness"`
	TemplateEvenness  []TemplateEvennessItem `json:"template_evenness"`
	Violations        []string               `json:"violations"`
}

type CreationStats struct {
	ShiftsCreated      int      `json:"shifts_created"`
	AssignmentsCreated int      `json:"assignments_created"`
	UnfilledCount      int      `json:"unfilled_count"`
	Errors             []string `json:"errors"`
}

type PlanResponse struct {
	RunID         string           `json:"runId"`
	CreatedAtISO  string           `json:"createdAtISO"`
	Window        Window           `json:"window"`
	Assignments   []AssignmentOut  `json:"assignments"`
	Users         []UserSummaryOut `json:"users"`
	Fairness      FairnessOut      `json:"fairness"`
	InputPayload  PlanPayload      `json:"input_payload"`
	CreationStats CreationStats    `json:"creation_stats"`
}

type slot struct {
	Date              string
	TemplateID        int
	TemplatePriority  int
	Start             time.Time
	End               time.Time
	RoleID            int
	SlotIndex         int
	Required          bool
	PreAssignedUserID *int
}

type userState struct {
	UserID int
	Name   string
	Roles  map[int]bool
	Avail  []AvailabilitySlot

	Vacations []TimeOffPeriod
	SickDays  []TimeOffPeriod

	LastEnd *time.Time

	MinutesByDay  map[string]int
	MinutesByWeek map[string]int
	AssignedSlots []AssignmentOut

	TotalShifts  int
	TotalMinutes int

	TemplateCounts map[int]int

	WorkedWeekendDays map[string]bool
}

func parseISO(ts string) (time.Time, error) { return time.Parse(time.RFC3339, ts) }

func dayKey(t time.Time) string { return t.Format("2006-01-02") }

func weekStartMonday(t time.Time, loc *time.Location) time.Time {
	tt := t.In(loc)
	wd := int(tt.Weekday())
	if wd == 0 {
		wd = 7
	}
	d := tt.AddDate(0, 0, -(wd - 1))
	y, m, day := d.Date()
	return time.Date(y, m, day, 0, 0, 0, 0, loc)
}

func minutesBetween(a time.Time, b time.Time) int {
	return int(b.Sub(a).Minutes())
}

func hhmmToParts(hhmm string) (int, int) {
	if len(hhmm) != 5 {
		return 0, 0
	}
	h := int(hhmm[0]-'0')*10 + int(hhmm[1]-'0')
	m := int(hhmm[3]-'0')*10 + int(hhmm[4]-'0')
	return h, m
}

func mandatoryBreakMinutes(workMinutes int) int {
	if workMinutes > 9*60 {
		return 45
	}
	if workMinutes > 6*60 {
		return 30
	}
	return 0
}

func dayMaskMatches(day time.Time, mask int) bool {
	wd := int(day.Weekday()) // 0=Sun..6=Sat
	var bit int
	if wd == 0 {
		bit = 64
	} else {
		bit = 1 << (wd - 1) // Mon=1, Tue=2, Wed=4, Thu=8, Fri=16, Sat=32
	}
	return (mask & bit) != 0
}

func withinAvailabilitySpan(start time.Time, end time.Time, slots []AvailabilitySlot, enforce bool) bool {
	if !enforce || len(slots) == 0 {
		return true
	}
	if dayKey(start) != dayKey(end) {
		return false
	}
	dow := int(start.Weekday())
	ds := start.Format("15:04")
	es := end.Format("15:04")
	for _, s := range slots {
		if s.DayOfWeek != dow {
			continue
		}
		if ds >= s.Start && es <= s.End {
			return true
		}
	}
	return false
}

func isOnTimeOff(start time.Time, end time.Time, periods []TimeOffPeriod) bool {
	for _, p := range periods {
		ps, err1 := parseISO(p.FromISO)
		pe, err2 := parseISO(p.ToISO)
		if err1 != nil || err2 != nil {
			continue
		}
		if start.Before(pe) && end.After(ps) {
			return true
		}
	}
	return false
}

func overlapsAny(existing []AssignmentOut, start time.Time, end time.Time) bool {
	for _, a := range existing {
		s, err1 := parseISO(a.StartISO)
		e, err2 := parseISO(a.EndISO)
		if err1 != nil || err2 != nil {
			continue
		}
		if end.After(s) && start.Before(e) {
			return true
		}
	}
	return false
}

func hasRest(lastEnd *time.Time, nextStart time.Time, minRestHours int) bool {
	if lastEnd == nil {
		return true
	}
	le := *lastEnd
	if nextStart.Sub(le) >= time.Duration(minRestHours)*time.Hour {
		return true
	}
	return false
}

func startOfDay(t time.Time, loc *time.Location) time.Time {
	tt := t.In(loc)
	y, m, d := tt.Date()
	return time.Date(y, m, d, 0, 0, 0, 0, loc)
}

func isWeekendDay(t time.Time) bool {
	wd := int(t.Weekday())
	return wd == 6 || wd == 0
}

func weekendKeyForDate(t time.Time, loc *time.Location) string {
	tt := t.In(loc)
	ws := weekStartMonday(tt, loc)
	sat := ws.AddDate(0, 0, 5)
	return sat.Format("2006-01-02")
}

func computeFreeWeekendsByUser(u *userState, loc *time.Location, windowStart time.Time, windowEnd time.Time) (int, int) {
	weekendsWorked := map[string]bool{}
	for _, a := range u.AssignedSlots {
		s, err := parseISO(a.StartISO)
		if err != nil {
			continue
		}
		if isWeekendDay(s.In(loc)) {
			wk := weekendKeyForDate(s, loc)
			weekendsWorked[wk] = true
		}
	}
	totalWeekends := 0
	for d := startOfDay(windowStart, loc); !d.After(windowEnd); d = d.AddDate(0, 0, 1) {
		if int(d.Weekday()) == 6 {
			totalWeekends++
		}
	}
	worked := len(weekendsWorked)
	free := totalWeekends - worked
	if free < 0 {
		free = 0
	}
	return free, worked
}

func buildSlots(payload PlanPayload, loc *time.Location, windowStart time.Time, windowEnd time.Time) ([]slot, int) {
	slots := make([]slot, 0, 1024)
	shiftsCreated := 0

	// Build holiday data map for quick lookup
	holidayData := make(map[int]HolidayInput)
	for _, h := range payload.Holidays {
		holidayData[h.ID] = h
	}

	// Build existing shifts map for quick lookup (key: "YYYY-MM-DD_templateID_HH:MM")
	existingShifts := make(map[string]bool)
	for _, es := range payload.ExistingShifts {
		key := es.Date + "_" + itoa(es.TemplateID) + "_" + es.StartTime
		existingShifts[key] = true
	}

	type roleSlotDef struct {
		roleID   int
		required bool
		pre      *int
	}

	for day := startOfDay(windowStart, loc); !day.After(windowEnd); day = day.AddDate(0, 0, 1) {
		// Determine which templates are overridden by holiday templates on this day
		overriddenParentIDs := make(map[int]bool)
		for _, t := range payload.ShiftTemplates {
			if t.HolidayMode && t.ParentTemplateID != nil {
				// This is a holiday override template - check if it applies to this day
				if isHolidayDate(day, t.HolidayIds, holidayData) {
					overriddenParentIDs[*t.ParentTemplateID] = true
				}
			}
		}

		for _, t := range payload.ShiftTemplates {
			// Skip parent templates on days when they are overridden by holiday templates
			if !t.HolidayMode && overriddenParentIDs[t.ID] {
				continue
			}

			// Holiday override templates only apply on their specific holiday dates
			if t.HolidayMode && t.ParentTemplateID != nil {
				if !isHolidayDate(day, t.HolidayIds, holidayData) {
					continue
				}
			}

			// Regular templates use day mask
			if !t.HolidayMode && !dayMaskMatches(day, t.DayMask) {
				continue
			}

			// Check if shift already exists for this date and template
			dayStr := day.Format("2006-01-02")
			existingKey := dayStr + "_" + itoa(t.ID) + "_" + t.StartTime
			if existingShifts[existingKey] {
				// Skip creating this shift - it already exists
				continue
			}

			sh, sm := hhmmToParts(t.StartTime)
			eh, em := hhmmToParts(t.EndTime)
			start := time.Date(day.Year(), day.Month(), day.Day(), sh, sm, 0, 0, loc)
			end := time.Date(day.Year(), day.Month(), day.Day(), eh, em, 0, 0, loc)

			// Handle overnight shifts: if end <= start, shift goes to next day
			if !end.After(start) {
				end = end.AddDate(0, 0, 1)
			}

			roleSlots := make([]roleSlotDef, 0, 16)
			for _, r := range t.MustHaveRoles {
				count := r.Count
				if count <= 0 {
					continue
				}
				for i := 0; i < count; i++ {
					var pre *int
					if i < len(r.PreAssignedUsers) && r.PreAssignedUsers[i] != nil {
						pre = r.PreAssignedUsers[i]
					}
					roleSlots = append(roleSlots, roleSlotDef{roleID: r.RoleID, required: true, pre: pre})
				}
			}
			for _, r := range t.NiceToHaveRoles {
				count := r.Count
				if count <= 0 {
					continue
				}
				for i := 0; i < count; i++ {
					var pre *int
					if i < len(r.PreAssignedUsers) && r.PreAssignedUsers[i] != nil {
						pre = r.PreAssignedUsers[i]
					}
					roleSlots = append(roleSlots, roleSlotDef{roleID: r.RoleID, required: false, pre: pre})
				}
			}

			perRoleIndex := map[int]int{}
			for _, rs := range roleSlots {
				idx := perRoleIndex[rs.roleID]
				perRoleIndex[rs.roleID] = idx + 1
				slots = append(slots, slot{
					Date:              day.Format("2006-01-02"),
					TemplateID:        t.ID,
					TemplatePriority:  t.Priority,
					Start:             start,
					End:               end,
					RoleID:            rs.roleID,
					SlotIndex:         idx,
					Required:          rs.required,
					PreAssignedUserID: rs.pre,
				})
				shiftsCreated++
			}
		}
	}

	sort.SliceStable(slots, func(i, j int) bool {
		if slots[i].Start.Equal(slots[j].Start) {
			if slots[i].TemplatePriority != slots[j].TemplatePriority {
				return slots[i].TemplatePriority > slots[j].TemplatePriority
			}
			if slots[i].TemplateID != slots[j].TemplateID {
				return slots[i].TemplateID < slots[j].TemplateID
			}
			if slots[i].RoleID != slots[j].RoleID {
				return slots[i].RoleID < slots[j].RoleID
			}
			return slots[i].SlotIndex < slots[j].SlotIndex
		}
		return slots[i].Start.Before(slots[j].Start)
	})

	return slots, shiftsCreated
}

func computeTargets(slots []slot, userCount int) (map[int]float64, float64) {
	templateTotals := map[int]int{}
	requiredTotal := 0
	for _, s := range slots {
		if s.Required {
			templateTotals[s.TemplateID]++
			requiredTotal++
		}
	}
	templateTargetPerUser := map[int]float64{}
	for tid, c := range templateTotals {
		templateTargetPerUser[tid] = float64(c) / float64(userCount)
	}
	totalTarget := float64(requiredTotal) / float64(userCount)
	return templateTargetPerUser, totalTarget
}

func eligible(u *userState, sl slot, rules RulesInput, loc *time.Location) (int, bool) {
	if !u.Roles[sl.RoleID] {
		return 0, false
	}
	if isOnTimeOff(sl.Start, sl.End, u.Vacations) || isOnTimeOff(sl.Start, sl.End, u.SickDays) {
		return 0, false
	}
	if !withinAvailabilitySpan(sl.Start, sl.End, u.Avail, rules.EnforceAvailability) {
		return 0, false
	}
	if overlapsAny(u.AssignedSlots, sl.Start, sl.End) {
		return 0, false
	}
	if !hasRest(u.LastEnd, sl.Start, rules.MinRestHours) {
		return 0, false
	}

	total := minutesBetween(sl.Start, sl.End)
	if total <= 0 {
		return 0, false
	}
	paid := total - mandatoryBreakMinutes(total)
	if paid <= 0 {
		return 0, false
	}

	if rules.MaxDailyHours > 0 {
		dk := dayKey(sl.Start.In(loc))
		if u.MinutesByDay[dk]+paid > rules.MaxDailyHours*60 {
			return 0, false
		}
	}
	if rules.MaxHoursPerWeek > 0 {
		wk := weekStartMonday(sl.Start, loc).Format("2006-01-02")
		if u.MinutesByWeek[wk]+paid > rules.MaxHoursPerWeek*60 {
			return 0, false
		}
	}

	if rules.MaxConsecutiveDays > 0 {
		slotDay := startOfDay(sl.Start, loc)
		hasPrev := false
		var prevDay time.Time
		if u.LastEnd != nil {
			prevDay = startOfDay((*u.LastEnd).In(loc), loc)
			hasPrev = true
		}
		if hasPrev {
			daysDiff := int(slotDay.Sub(prevDay).Hours() / 24)
			if daysDiff == 0 {
				return paid, true
			}
			if daysDiff == 1 {
				cons := 1
				for back := 1; back <= rules.MaxConsecutiveDays; back++ {
					dk := dayKey(slotDay.AddDate(0, 0, -back).In(loc))
					if u.MinutesByDay[dk] > 0 {
						cons++
					} else {
						break
					}
				}
				if cons > rules.MaxConsecutiveDays {
					return 0, false
				}
			}
		}
	}

	return paid, true
}

func chooseUser(sl slot, users []*userState, rules RulesInput, loc *time.Location, templateTarget map[int]float64, totalTarget float64, fairness FairnessInput, windowStart time.Time, windowEnd time.Time) (*userState, int) {
	if sl.PreAssignedUserID != nil {
		pre := *sl.PreAssignedUserID
		for _, u := range users {
			if u.UserID == pre {
				paid, ok := eligible(u, sl, rules, loc)
				if !ok {
					return nil, 0
				}
				return u, paid
			}
		}
		return nil, 0
	}

	type cand struct {
		u     *userState
		paid  int
		score float64
	}
	cands := make([]cand, 0, len(users))

	for _, u := range users {
		paid, ok := eligible(u, sl, rules, loc)
		if !ok {
			continue
		}

		// 1. Total shift balance - prefer users with fewer total shifts
		afterTotal := float64(u.TotalShifts + 1)
		totalPenalty := math.Abs(afterTotal - totalTarget)

		// 2. Template balance - prefer users with fewer shifts of this template type
		afterTemplate := float64(u.TemplateCounts[sl.TemplateID] + 1)
		templatePenalty := 0.0
		if tgt, ok := templateTarget[sl.TemplateID]; ok {
			templatePenalty = math.Abs(afterTemplate - tgt)
		}

		// 3. Weekend fairness - heavily penalize assigning weekends to users who already worked weekends
		weekendPenalty := 0.0
		if isWeekendDay(sl.Start.In(loc)) {
			free, worked := computeFreeWeekendsByUser(u, loc, windowStart, windowEnd)
			// Strong penalty if user has fewer free weekends than minimum
			if free < fairness.MinFreeWeekends {
				weekendPenalty = float64(fairness.MinFreeWeekends-free+1) * 2.0
			} else {
				// Prefer users who worked fewer weekends
				weekendPenalty = float64(worked) * 0.5
			}
		}

		// 4. Weekly load balancing - prefer users with fewer minutes this week
		wkKey := weekStartMonday(sl.Start, loc).Format("2006-01-02")
		weekMinutes := float64(u.MinutesByWeek[wkKey])
		weekPenalty := weekMinutes / 1500.0 // Normalize by ~38.5 hours

		// 5. Consecutive days penalty - prefer spreading work across different days
		consecutivePenalty := 0.0
		if u.LastEnd != nil {
			slotDay := startOfDay(sl.Start, loc)
			prevDay := startOfDay((*u.LastEnd).In(loc), loc)
			daysDiff := int(slotDay.Sub(prevDay).Hours() / 24)

			// Penalty for consecutive days (prefer some rest between shifts)
			if daysDiff == 1 {
				// Count consecutive days
				cons := 1
				for back := 1; back <= 7; back++ {
					dk := dayKey(slotDay.AddDate(0, 0, -back).In(loc))
					if u.MinutesByDay[dk] > 0 {
						cons++
					} else {
						break
					}
				}
				// Increasing penalty for more consecutive days
				if cons >= 3 {
					consecutivePenalty = float64(cons-2) * 0.5
				}
			}
		}

		// Weighted score - lower is better
		// Weights prioritize: weekend fairness > total balance > consecutive days > template balance > weekly load
		score := 2.5*totalPenalty + 1.5*templatePenalty + 4.0*weekendPenalty + 0.8*weekPenalty + 1.2*consecutivePenalty
		cands = append(cands, cand{u: u, paid: paid, score: score})
	}

	if len(cands) == 0 {
		return nil, 0
	}

	sort.SliceStable(cands, func(i, j int) bool {
		if math.Abs(cands[i].score-cands[j].score) > 1e-9 {
			return cands[i].score < cands[j].score
		}
		if cands[i].u.TotalShifts != cands[j].u.TotalShifts {
			return cands[i].u.TotalShifts < cands[j].u.TotalShifts
		}
		return cands[i].u.UserID < cands[j].u.UserID
	})

	return cands[0].u, cands[0].paid
}

func assignSlot(u *userState, sl slot, paid int, loc *time.Location) AssignmentOut {
	var uid int = u.UserID
	out := AssignmentOut{
		Date:       sl.Date,
		TemplateID: sl.TemplateID,
		StartISO:   sl.Start.Format(time.RFC3339),
		EndISO:     sl.End.Format(time.RFC3339),
		RoleID:     sl.RoleID,
		SlotIndex:  sl.SlotIndex,
		UserID:     &uid,
	}

	u.AssignedSlots = append(u.AssignedSlots, out)
	u.TotalShifts++
	u.TotalMinutes += paid
	u.TemplateCounts[sl.TemplateID]++

	dk := dayKey(sl.Start.In(loc))
	u.MinutesByDay[dk] += paid
	wk := weekStartMonday(sl.Start, loc).Format("2006-01-02")
	u.MinutesByWeek[wk] += paid

	le := sl.End
	u.LastEnd = &le

	if isWeekendDay(sl.Start.In(loc)) {
		u.WorkedWeekendDays[dk] = true
	}

	return out
}

func computeFairness(users []*userState, allSlots []slot, tol float64, fairness FairnessInput, loc *time.Location, windowStart time.Time, windowEnd time.Time) FairnessOut {
	activeUsers := make([]*userState, 0, len(users))
	for _, u := range users {
		activeUsers = append(activeUsers, u)
	}
	if len(activeUsers) == 0 {
		return FairnessOut{TolerancePct: tol, OverallScore: 0, Violations: []string{"no_users"}}
	}

	totalCounts := make([]int, 0, len(activeUsers))
	minTotal := int(^uint(0) >> 1)
	maxTotal := 0
	sumTotal := 0
	for _, u := range activeUsers {
		totalCounts = append(totalCounts, u.TotalShifts)
		sumTotal += u.TotalShifts
		if u.TotalShifts < minTotal {
			minTotal = u.TotalShifts
		}
		if u.TotalShifts > maxTotal {
			maxTotal = u.TotalShifts
		}
	}
	avgTotal := float64(sumTotal) / float64(len(activeUsers))
	maxDiffPct := 0.0
	if avgTotal > 0 {
		maxDiffPct = float64(maxTotal-minTotal) / avgTotal
	}

	templateTotals := map[int][]int{}
	templateSet := map[int]bool{}
	for _, sl := range allSlots {
		if sl.Required {
			templateSet[sl.TemplateID] = true
		}
	}
	templateIDs := make([]int, 0, len(templateSet))
	for tid := range templateSet {
		templateIDs = append(templateIDs, tid)
	}
	sort.Ints(templateIDs)

	for _, tid := range templateIDs {
		templateTotals[tid] = make([]int, 0, len(activeUsers))
		for _, u := range activeUsers {
			templateTotals[tid] = append(templateTotals[tid], u.TemplateCounts[tid])
		}
	}

	templateEvenness := make([]TemplateEvennessItem, 0, len(templateIDs))
	for _, tid := range templateIDs {
		vals := templateTotals[tid]
		minV := int(^uint(0) >> 1)
		maxV := 0
		sumV := 0
		for _, v := range vals {
			sumV += v
			if v < minV {
				minV = v
			}
			if v > maxV {
				maxV = v
			}
		}
		avgV := float64(sumV) / float64(len(vals))
		maxDiff := float64(maxV) - avgV
		if avgV > float64(maxV) {
			maxDiff = avgV - float64(minV)
		}
		if avgV > 0 {
			if (avgV - float64(minV)) > (float64(maxV) - avgV) {
				maxDiff = avgV - float64(minV)
			} else {
				maxDiff = float64(maxV) - avgV
			}
		}
		templateEvenness = append(templateEvenness, TemplateEvennessItem{
			TemplateID:     tid,
			Avg:            avgV,
			Min:            minV,
			Max:            maxV,
			MaxDiffFromAvg: maxDiff,
		})
	}

	violations := make([]string, 0, 8)
	if maxDiffPct > tol {
		violations = append(violations, "total_shifts_imbalance_gt_tol")
	}
	for _, u := range activeUsers {
		free, _ := computeFreeWeekendsByUser(u, loc, windowStart, windowEnd)
		if free < fairness.MinFreeWeekends {
			violations = append(violations, "min_free_weekends_not_met: user "+itoa(u.UserID))
		}
	}

	score := 1.0
	score -= clamp01(maxDiffPct) * 0.6
	templatePenalty := 0.0
	for _, it := range templateEvenness {
		if it.Avg <= 0 {
			continue
		}
		templatePenalty += clamp01(it.MaxDiffFromAvg / it.Avg)
	}
	if len(templateEvenness) > 0 {
		templatePenalty /= float64(len(templateEvenness))
	}
	score -= templatePenalty * 0.3
	freePenalty := 0.0
	for _, u := range activeUsers {
		free, _ := computeFreeWeekendsByUser(u, loc, windowStart, windowEnd)
		if free < fairness.MinFreeWeekends {
			freePenalty += float64(fairness.MinFreeWeekends-free) / float64(fairness.MinFreeWeekends)
		}
	}
	if len(activeUsers) > 0 {
		freePenalty /= float64(len(activeUsers))
	}
	score -= clamp01(freePenalty) * 0.2

	score = clamp01(score)

	return FairnessOut{
		TolerancePct: tol,
		OverallScore: score,
		UserTotalEvenness: TotalEvenness{
			Avg:        avgTotal,
			Min:        minTotal,
			Max:        maxTotal,
			MaxDiffPct: maxDiffPct,
		},
		TemplateEvenness: templateEvenness,
		Violations:       violations,
	}
}

func clamp01(v float64) float64 {
	if v < 0 {
		return 0
	}
	if v > 1 {
		return 1
	}
	return v
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	neg := false
	if n < 0 {
		neg = true
		n = -n
	}
	buf := make([]byte, 0, 16)
	for n > 0 {
		d := n % 10
		buf = append(buf, byte('0'+d))
		n /= 10
	}
	if neg {
		buf = append(buf, '-')
	}
	for i, j := 0, len(buf)-1; i < j; i, j = i+1, j-1 {
		buf[i], buf[j] = buf[j], buf[i]
	}
	return string(buf)
}

// calculateEaster calculates Easter Sunday for a given year using the Computus algorithm (Gauss formula).
func calculateEaster(year int) time.Time {
	a := year % 19
	b := year / 100
	c := year % 100
	d := b / 4
	e := b % 4
	f := (b + 8) / 25
	g := (b - f + 1) / 3
	h := (19*a + b - d - g + 15) % 30
	i := c / 4
	k := c % 4
	l := (32 + 2*e + 2*i - h - k) % 7
	m := (a + 11*h + 22*l) / 451
	month := (h + l - 7*m + 114) / 31
	day := ((h + l - 7*m + 114) % 31) + 1

	return time.Date(year, time.Month(month), day, 0, 0, 0, 0, time.UTC)
}

// isHolidayDate checks if a given date matches any of the specified holiday IDs
func isHolidayDate(date time.Time, holidayIds []int, holidayData map[int]HolidayInput) bool {
	if len(holidayIds) == 0 || len(holidayData) == 0 {
		return false
	}

	for _, hid := range holidayIds {
		holiday, ok := holidayData[hid]
		if !ok {
			continue
		}

		var holidayDate time.Time
		if holiday.CalculationType == "easter_relative" {
			easter := calculateEaster(date.Year())
			holidayDate = easter.AddDate(0, 0, holiday.EasterOffset)
		} else {
			// Fixed date - check for leap year edge case
			if holiday.Month == 2 && holiday.Day == 29 && !isLeapYear(date.Year()) {
				continue
			}
			holidayDate = time.Date(date.Year(), time.Month(holiday.Month), holiday.Day, 0, 0, 0, 0, date.Location())
		}

		if isSameDay(date, holidayDate) {
			return true
		}
	}

	return false
}

// isSameDay checks if two times represent the same calendar day
func isSameDay(a, b time.Time) bool {
	ay, am, ad := a.Date()
	by, bm, bd := b.Date()
	return ay == by && am == bm && ad == bd
}

// isLeapYear checks if a year is a leap year
func isLeapYear(year int) bool {
	return year%4 == 0 && (year%100 != 0 || year%400 == 0)
}

func deepCopyUsers(base []*userState) []*userState {
	out := make([]*userState, 0, len(base))
	for _, b := range base {
		u := &userState{
			UserID: b.UserID,
			Name:   b.Name,
			Roles:  map[int]bool{},
			Avail:  append([]AvailabilitySlot{}, b.Avail...),

			Vacations: append([]TimeOffPeriod{}, b.Vacations...),
			SickDays:  append([]TimeOffPeriod{}, b.SickDays...),

			LastEnd: nil,

			MinutesByDay:  map[string]int{},
			MinutesByWeek: map[string]int{},
			AssignedSlots: make([]AssignmentOut, 0, 256),

			TotalShifts:  0,
			TotalMinutes: 0,

			TemplateCounts:    map[int]int{},
			WorkedWeekendDays: map[string]bool{},
		}
		for k, v := range b.Roles {
			u.Roles[k] = v
		}
		if b.LastEnd != nil {
			le := *b.LastEnd
			u.LastEnd = &le
		}
		out = append(out, u)
	}
	return out
}

func main() {
	var runIDFlag string
	flag.StringVar(&runIDFlag, "run-id", "", "")
	flag.Parse()

	dec := json.NewDecoder(os.Stdin)
	var req PlanRequest
	err := dec.Decode(&req)
	if err != nil && err != io.EOF {
		os.Exit(2)
	}

	payload := req.Payload
	if payload.Timezone == "" {
		payload.Timezone = "Europe/Berlin"
	}
	loc, locErr := time.LoadLocation(payload.Timezone)
	if locErr != nil {
		loc, _ = time.LoadLocation("Europe/Berlin")
	}

	if payload.Year <= 0 || payload.Month <= 0 || payload.Month > 12 {
		os.Exit(2)
	}

	windowStart := time.Date(payload.Year, time.Month(payload.Month), 1, 0, 0, 0, 0, loc)
	windowEnd := windowStart.AddDate(0, 1, 0).Add(-time.Second)

	baseUsers := make([]*userState, 0, len(payload.Users))
	for _, ui := range payload.Users {
		roleMap := map[int]bool{}
		for _, r := range ui.Roles {
			roleMap[r] = true
		}
		var lastEnd *time.Time
		if ui.LastEndISO != nil && *ui.LastEndISO != "" {
			t, e := parseISO(*ui.LastEndISO)
			if e == nil {
				tt := t.In(loc)
				lastEnd = &tt
			}
		}
		u := &userState{
			UserID: ui.UserID,
			Name:   ui.Name,
			Roles:  roleMap,
			Avail:  append([]AvailabilitySlot{}, ui.Availability...),

			Vacations: append([]TimeOffPeriod{}, ui.Vacations...),
			SickDays:  append([]TimeOffPeriod{}, ui.SickDays...),

			LastEnd: lastEnd,

			MinutesByDay:  map[string]int{},
			MinutesByWeek: map[string]int{},
			AssignedSlots: make([]AssignmentOut, 0, 256),

			TotalShifts:  0,
			TotalMinutes: 0,

			TemplateCounts:    map[int]int{},
			WorkedWeekendDays: map[string]bool{},
		}
		baseUsers = append(baseUsers, u)
	}

	allSlots, shiftsCreated := buildSlots(payload, loc, windowStart, windowEnd)

	maxIterations := 10
	bestAssignments := make([]AssignmentOut, 0, 1024)
	bestUsers := make([]*userState, 0, len(baseUsers))
	bestUnfilled := int(^uint(0) >> 1)
	bestScore := -1.0

	templateTarget, totalTarget := computeTargets(allSlots, len(baseUsers))
	tol := payload.Fairness.TotalShiftTolerancePct
	if tol <= 0 {
		tol = 0.1
	}

	for iter := 1; iter <= maxIterations; iter++ {
		iterUsers := deepCopyUsers(baseUsers)
		iterAssignments := make([]AssignmentOut, 0, len(allSlots))
		unfilled := 0

		requiredSlots := make([]slot, 0, len(allSlots))
		optionalSlots := make([]slot, 0, len(allSlots))
		for _, s := range allSlots {
			if s.Required {
				requiredSlots = append(requiredSlots, s)
			} else {
				optionalSlots = append(optionalSlots, s)
			}
		}

		for _, sl := range requiredSlots {
			u, paid := chooseUser(sl, iterUsers, payload.Rules, loc, templateTarget, totalTarget, payload.Fairness, windowStart, windowEnd)
			if u == nil {
				note := "UNFILLED"
				iterAssignments = append(iterAssignments, AssignmentOut{
					Date:       sl.Date,
					TemplateID: sl.TemplateID,
					StartISO:   sl.Start.Format(time.RFC3339),
					EndISO:     sl.End.Format(time.RFC3339),
					RoleID:     sl.RoleID,
					SlotIndex:  sl.SlotIndex,
					Note:       &note,
				})
				unfilled++
				continue
			}
			out := assignSlot(u, sl, paid, loc)
			iterAssignments = append(iterAssignments, out)
		}

		for _, sl := range optionalSlots {
			u, paid := chooseUser(sl, iterUsers, payload.Rules, loc, templateTarget, totalTarget, payload.Fairness, windowStart, windowEnd)
			if u == nil {
				continue
			}
			out := assignSlot(u, sl, paid, loc)
			iterAssignments = append(iterAssignments, out)
		}

		f := computeFairness(iterUsers, allSlots, tol, payload.Fairness, loc, windowStart, windowEnd)
		score := f.OverallScore

		isBetter := false
		if unfilled < bestUnfilled {
			isBetter = true
		} else if unfilled == bestUnfilled && score > bestScore {
			isBetter = true
		}
		if isBetter || iter == 1 {
			bestUnfilled = unfilled
			bestScore = score
			bestAssignments = iterAssignments
			bestUsers = iterUsers
		}
	}

	userSummaries := make([]UserSummaryOut, 0, len(bestUsers))
	for _, u := range bestUsers {
		tc := map[string]int{}
		for tid, c := range u.TemplateCounts {
			tc[itoa(tid)] = c
		}
		tp := map[string]float64{}
		if u.TotalShifts > 0 {
			for tid, c := range u.TemplateCounts {
				tp[itoa(tid)] = float64(c) / float64(u.TotalShifts)
			}
		}
		freeW, workedW := computeFreeWeekendsByUser(u, loc, windowStart, windowEnd)
		userSummaries = append(userSummaries, UserSummaryOut{
			UserID:          u.UserID,
			Name:            u.Name,
			TotalShifts:     u.TotalShifts,
			TotalMinutes:    u.TotalMinutes,
			TemplateCounts:  tc,
			TemplatePercent: tp,
			FreeWeekends:    freeW,
			WeekendWorked:   workedW,
		})
	}
	sort.SliceStable(userSummaries, func(i, j int) bool { return userSummaries[i].UserID < userSummaries[j].UserID })

	fout := computeFairness(bestUsers, allSlots, tol, payload.Fairness, loc, windowStart, windowEnd)

	assignedCount := 0
	for _, a := range bestAssignments {
		if a.UserID != nil {
			assignedCount++
		}
	}

	resp := PlanResponse{
		RunID: func() string {
			if runIDFlag != "" {
				return runIDFlag
			}
			return "plan"
		}(),
		CreatedAtISO: time.Now().UTC().Format(time.RFC3339),
		Window: Window{
			StartISO: windowStart.Format(time.RFC3339),
			EndISO:   windowEnd.Format(time.RFC3339),
		},
		Assignments:  bestAssignments,
		Users:        userSummaries,
		Fairness:     fout,
		InputPayload: payload,
		CreationStats: CreationStats{
			ShiftsCreated:      shiftsCreated,
			AssignmentsCreated: assignedCount,
			UnfilledCount:      bestUnfilled,
			Errors:             []string{},
		},
	}

	enc := json.NewEncoder(os.Stdout)
	enc.SetEscapeHTML(false)
	_ = enc.Encode(resp)
}
