package main

import (
	"context"
	"fmt"
	"time"

	"github.com/redis/go-redis/v9"
)

// AtomicLuaReserve evaluates available stock, decrements counter, and sets TTL reservation atomically.
const AtomicLuaReserve = `
local stock_key = KEYS[1]
local res_key   = KEYS[2]
local req_qty   = tonumber(ARGV[1])
local ttl_sec   = tonumber(ARGV[2])

local current_stock = tonumber(redis.call("GET", stock_key) or "0")
if current_stock >= req_qty then
    redis.call("DECRBY", stock_key, req_qty)
    redis.call("SETEX", res_key, ttl_sec, req_qty)
    return 1
else
    return 0
end
`

type InventoryEngine struct {
	client *redis.Client
	script *redis.Script
}

// NewInventoryEngine initializes the Redis client and pre-compiles the Lua script SHA.
func NewInventoryEngine(client *redis.Client) *InventoryEngine {
	return &InventoryEngine{
		client: client,
		script: redis.NewScript(AtomicLuaReserve),
	}
}

// ReserveATP executes the atomic soft reservation for an order SKU.
func (e *InventoryEngine) ReserveATP(ctx context.Context, warehouseID, sku, orderID string, qty int, ttl time.Duration) (bool, error) {
	stockKey := fmt.Sprintf("atp:{%s:%s}:stock", warehouseID, sku)
	resKey := fmt.Sprintf("atp:{%s:%s}:res:%s", warehouseID, sku, orderID)

	keys := []string{stockKey, resKey}
	args := []interface{}{qty, int(ttl.Seconds())}

	res, err := e.script.Run(ctx, e.client, keys, args...).Int64()
	if err != nil {
		return false, fmt.Errorf("redis ATP reservation script failed: %w", err)
	}

	return res == 1, nil
}

// File: internal/fulfillment/allocator/solver.go
package allocator

import (
	"context"
	"fmt"
	"math"
)

// GeoLocation represents latitude and longitude coordinates.
type GeoLocation struct {
	Lat float64
	Lon float64
}

// Warehouse models a fulfillment center facility.
type Warehouse struct {
	ID        string
	Location  GeoLocation
	LaborCost float64
	Inventory map[string]int // SKU -> Available Quantity
}

// CartItem represents an item requested by the customer.
type CartItem struct {
	SKU      string
	Quantity int
}

// ShipmentPackage represents a group of items shipped from one warehouse.
type ShipmentPackage struct {
	WarehouseID string
	Items       []CartItem
	FreightCost float64
}

// AllocationPlan details the complete fulfillment strategy for an order.
type AllocationPlan struct {
	Packages    []ShipmentPackage
	TotalCost   float64
	SplitCount  int
}

// Solver executes multi-warehouse constraint optimization.
type Solver struct {
	warehouses    []Warehouse
	splitPenalty  float64
	costPerKm     float64
}

func NewSolver(warehouses []Warehouse, splitPenalty, costPerKm float64) *Solver {
	return &Solver{
		warehouses:   warehouses,
		splitPenalty: splitPenalty,
		costPerKm:    costPerKm,
	}
}

// Haversine computes great-circle distance between two geographic coordinates in kilometers.
func haversine(p1, p2 GeoLocation) float64 {
	const earthRadiusKm = 6371.0
	dLat := (p2.Lat - p1.Lat) * (math.Pi / 180.0)
	dLon := (p2.Lon - p1.Lon) * (math.Pi / 180.0)

	lat1 := p1.Lat * (math.Pi / 180.0)
	lat2 := p2.Lat * (math.Pi / 180.0)

	a := math.Sin(dLat/2)*math.Sin(dLat/2) +
		math.Sin(dLon/2)*math.Sin(dLon/2)*math.Cos(lat1)*math.Cos(lat2)
	c := 2 * math.Atan2(math.Sqrt(a), math.Sqrt(1-a))

	return earthRadiusKm * c
}

// Solve determines the optimal warehouse assignment minimizing freight, labor, and split penalties.
func (s *Solver) Solve(ctx context.Context, customerLoc GeoLocation, cart []CartItem) (*AllocationPlan, error) {
	// 1. Check if a single warehouse can fulfill 100% of the cart (Zero-Split Preferred)
	var bestSingleWH *Warehouse
	var minSingleCost = math.MaxFloat64

	for i := range s.warehouses {
		wh := &s.warehouses[i]
		canFulfillAll := true
		for _, item := range cart {
			if wh.Inventory[item.SKU] < item.Quantity {
				canFulfillAll = false
				break
			}
		}

		if canFulfillAll {
			dist := haversine(wh.Location, customerLoc)
			cost := (dist * s.costPerKm) + wh.LaborCost
			if cost < minSingleCost {
				minSingleCost = cost
				bestSingleWH = wh
			}
		}
	}

	if bestSingleWH != nil {
		dist := haversine(bestSingleWH.Location, customerLoc)
		return &AllocationPlan{
			Packages: []ShipmentPackage{
				{
					WarehouseID: bestSingleWH.ID,
					Items:       cart,
					FreightCost: dist * s.costPerKm,
				},
			},
			TotalCost:  minSingleCost,
			SplitCount: 0,
		}, nil
	}

	// 2. Multi-Warehouse Greedy Fallback (Minimize split penalty)
	unfulfilled := make(map[string]int)
	for _, item := range cart {
		unfulfilled[item.SKU] = item.Quantity
	}

	plan := &AllocationPlan{}

	for len(unfulfilled) > 0 {
		var selectedWH *Warehouse
		var maxCovered int
		var bestScore = math.MaxFloat64

		for i := range s.warehouses {
			wh := &s.warehouses[i]
			coveredCount := 0
			for sku, needed := range unfulfilled {
				if avail := wh.Inventory[sku]; avail > 0 {
					if avail >= needed {
						coveredCount += needed
					} else {
						coveredCount += avail
					}
				}
			}

			if coveredCount > 0 {
				dist := haversine(wh.Location, customerLoc)
				score := (dist * s.costPerKm) / float64(coveredCount)
				if score < bestScore {
					bestScore = score
					selectedWH = wh
					maxCovered = coveredCount
				}
			}
		}

		if selectedWH == nil || maxCovered == 0 {
			return nil, fmt.Errorf("insufficient global inventory to fulfill cart")
		}

		// Allocate items to selected warehouse
		pkgItems := make([]CartItem, 0)
		for sku, needed := range unfulfilled {
			avail := selectedWH.Inventory[sku]
			if avail > 0 {
				qtyToTake := needed
				if avail < needed {
					qtyToTake = avail
				}
				pkgItems = append(pkgItems, CartItem{SKU: sku, Quantity: qtyToTake})
				unfulfilled[sku] -= qtyToTake
				if unfulfilled[sku] == 0 {
					delete(unfulfilled, sku)
				}
			}
		}

		dist := haversine(selectedWH.Location, customerLoc)
		freight := dist * s.costPerKm
		plan.Packages = append(plan.Packages, ShipmentPackage{
			WarehouseID: selectedWH.ID,
			Items:       pkgItems,
			FreightCost: freight,
		})
		plan.TotalCost += freight + selectedWH.LaborCost
	}

	plan.SplitCount = len(plan.Packages) - 1
	plan.TotalCost += float64(plan.SplitCount) * s.splitPenalty

	return plan, nil
}

from typing import Dict, List, Tuple
from ortools.constraint_solver import routing_enums_pb2, pywrapcp

def solve_capacitated_vrp(
    distance_matrix: List[List[int]],
    demands: List[int],
    vehicle_capacities: List[int],
    depot: int = 0,
    time_limit_sec: int = 15,
) -> Tuple[Dict[int, List[int]], int]:
    """Solves Capacitated VRP with payload limits, drop penalties, and GLS metaheuristics."""
    num_locations = len(distance_matrix)
    num_vehicles = len(vehicle_capacities)

    manager = pywrapcp.RoutingIndexManager(num_locations, num_vehicles, depot)
    routing = pywrapcp.RoutingModel(manager)

    # 1. Transit Cost Callback
    def distance_callback(from_index: int, to_index: int) -> int:
        return distance_matrix[manager.IndexToNode(from_index)][manager.IndexToNode(to_index)]

    transit_idx = routing.RegisterTransitCallback(distance_callback)
    routing.SetArcCostEvaluatorOfAllVehicles(transit_idx)

    # 2. Vehicle Capacity Constraints
    def demand_callback(from_index: int) -> int:
        return demands[manager.IndexToNode(from_index)]

    demand_idx = routing.RegisterUnaryTransitCallback(demand_callback)
    routing.AddDimensionWithVehicleCapacity(
        demand_idx,
        0,                  # Null slack
        vehicle_capacities, # Vehicle capacity limits
        True,               # Cumulative starts at zero
        "Capacity"
    )

    # 3. Disjunctions: Penalty for unserviced stops when fleet capacity is exceeded
    penalty = 100_000
    for node in range(1, num_locations):
        routing.AddDisjunction([manager.NodeToIndex(node)], penalty)

    # 4. Search Metaheuristics
    search_params = pywrapcp.DefaultRoutingSearchParameters()
    search_params.first_solution_strategy = (
        routing_enums_pb2.FirstSolutionStrategy.PATH_CHEAPEST_ARC
    )
    search_params.local_search_metaheuristic = (
        routing_enums_pb2.LocalSearchMetaheuristic.GUIDED_LOCAL_SEARCH
    )
    search_params.time_limit.seconds = time_limit_sec
    search_params.log_search = False

    solution = routing.SolveWithParameters(search_params)
    if not solution:
        return {}, -1

    routes: Dict[int, List[int]] = {}
    for vehicle_id in range(num_vehicles):
        index = routing.Start(vehicle_id)
        route = []
        while not routing.IsEnd(index):
            route.append(manager.IndexToNode(index))
            index = solution.Value(routing.NextVar(index))
        route.append(manager.IndexToNode(index))
        routes[vehicle_id] = route

    return routes, solution.ObjectiveValue()