doormile_backend/internal/assignment/ai_layer.go

package assignment

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"os"
	"strconv"
	"time"

	"doormile/constants"
	"doormile/db"
	"doormile/models"
	"doormile/utils"

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

// ─── Request / response types ────────────────────────────────────────────────

type aiCandidate struct {
	MilerID        int     `json:"miler_id"`
	DistanceKm     float64 `json:"distance_km"`
	Rating         float64 `json:"rating"`
	ActiveBookings int64   `json:"active_bookings"`
	OnTimeRate30d  float64 `json:"on_time_rate_30d"`
	CompletedToday int64   `json:"completed_today"`
	HubID          int     `json:"hub_id"`
	HubLoad        int64   `json:"hub_load"`
	HubCapacity    int     `json:"hub_capacity"`
}

type aiDecisionRequest struct {
	Booking    aiBookingInfo `json:"booking"`
	Candidates []aiCandidate `json:"candidates"`
	Context    aiRequestCtx  `json:"context"`
}

type aiBookingInfo struct {
	PickupLat    float64 `json:"pickup_lat"`
	PickupLon    float64 `json:"pickup_lon"`
	DeliveryLat  float64 `json:"delivery_lat"`
	DeliveryLon  float64 `json:"delivery_lon"`
	ItemCategory string  `json:"item_category"`
	Weight       float64 `json:"weight"`
	ServiceType  string  `json:"service_type"`
}

type aiRequestCtx struct {
	Hour   int    `json:"hour"`
	IsPeak bool   `json:"is_peak"`
	Zone   string `json:"zone"`
}

type aiDecisionResponse struct {
	ChosenMilerID   int    `json:"chosen_miler_id"`
	Escalate        bool   `json:"escalate"`
	AgentDecisionID uint64 `json:"agent_decision_id"`
	Reasoning       string `json:"reasoning"`
}

// ─── Main entry point ────────────────────────────────────────────────────────

// selectMilerWithAI collects all eligible milers from the GEORADIUS result, calls
// the AI decision engine, and returns the chosen miler plus the decision ID for
// audit. Falls back to the original distance/load/rating formula when the AI
// layer is unreachable or times out. Returns (nil, nil, false) when there are no
// eligible candidates or when the AI layer escalates the booking.
func selectMilerWithAI(booking *models.PickupBooking, nearby []redis.GeoLocation) (*milerCandidate, []*milerCandidate, *uint64, bool) {
	// Never offer an order back to a rider it was released from, who
	// rejected it or who cancelled it.
	nearby = withoutRiders(nearby, ridersToSkip(booking.Bookingid))

	pool, poolAI := collectEligibleCandidates(nearby)
	if len(pool) == 0 {
		return nil, nil, nil, false
	}

	// Balance: only the riders holding the fewest orders are considered. The
	// AI (or the fallback) chooses among them, so the split stays even however
	// many orders and riders there are. The full pool goes back to the caller
	// for the final re-check at commit time (finalizeChoice).
	candidates, aiCandidates := leastLoaded(pool, poolAI)

	decision, err := callDecisionEngine(booking, aiCandidates)
	if err != nil {
		utils.Warn("AI_LAYER_FALLBACK: decide-assignment unreachable, using legacy scoring",
			"booking_id", booking.Bookingid, "error", err)
		best := pickBestFromCandidates(candidates)
		return best, pool, nil, best != nil
	}

	utils.Info("Assignment: AI layer responded",
		"booking_id", booking.Bookingid,
		"escalate", decision.Escalate,
		"chosen_miler_id", decision.ChosenMilerID,
		"agent_decision_id", decision.AgentDecisionID,
		"reasoning", decision.Reasoning,
	)

	if decision.Escalate {
		utils.Warn("Assignment: AI layer escalated — skipping assignment this attempt",
			"booking_id", booking.Bookingid,
			"reasoning", decision.Reasoning,
		)
		return nil, nil, nil, false
	}

	var decisionID *uint64
	if decision.AgentDecisionID != 0 {
		id := decision.AgentDecisionID
		decisionID = &id
	}

	for _, c := range candidates {
		if c.profile.Userid == decision.ChosenMilerID {
			return c, pool, decisionID, true
		}
	}

	// AI returned an ID that is not among the least-loaded riders — fall back.
	utils.Warn("AI_LAYER_FALLBACK: chosen miler not in eligible set, using legacy scoring",
		"booking_id", booking.Bookingid,
		"chosen_miler_id", decision.ChosenMilerID,
	)
	best := pickBestFromCandidates(candidates)
	return best, pool, nil, best != nil
}

// leastLoaded keeps the candidates holding the fewest orders in hand (and
// the matching AI rows, which are parallel to them; ai may be nil).
func leastLoaded(cands []*milerCandidate, ai []aiCandidate) ([]*milerCandidate, []aiCandidate) {
	if len(cands) == 0 {
		return cands, ai
	}
	least := cands[0].activeBookings
	for _, c := range cands[1:] {
		if c.activeBookings < least {
			least = c.activeBookings
		}
	}
	var outC []*milerCandidate
	var outAI []aiCandidate
	for i, c := range cands {
		if c.activeBookings == least {
			outC = append(outC, c)
			if i < len(ai) {
				outAI = append(outAI, ai[i])
			}
		}
	}
	return outC, outAI
}

// ─── Candidate collection ────────────────────────────────────────────────────

// collectEligibleCandidates iterates the GEORADIUS result, applies eligibility
// filters (availability, active-booking cap), and fetches per-miler stats needed
// by the AI layer. Returns parallel slices so callers can use either.
func collectEligibleCandidates(nearby []redis.GeoLocation) ([]*milerCandidate, []aiCandidate) {
	var candidates []*milerCandidate
	var aiCandidates []aiCandidate

	for _, loc := range nearby {
		milerUserID, err := strconv.Atoi(loc.Name)
		if err != nil {
			utils.Warn("Assignment: skipping non-numeric GEO member", "name", loc.Name)
			continue
		}

		var profile models.MilerProfile
		if err := db.DB.Where("userid = ?", milerUserID).First(&profile).Error; err != nil {
			continue
		}

		// Carrying an order is not a reason to be skipped — maxActive below
		// is what decides how much one miler can hold. Testing for Available
		// here made that cap unreachable: a rider was eligible only while
		// idle, so activeCount was always 0 and the multi-stop round the cap
		// was written for could never be built.
		if !constants.MilerCanTakeWork(profile.Availabilitystatus) {
			continue
		}

		// Only a rider whose app is actually reporting can take an order. A
		// status left at "Assigned" by an app that stopped weeks ago used to
		// make that rider look like the least busy one, so orders went to
		// nobody. Compared in SQL against the database clock, which is right
		// whatever type the column has.
		if !milerHasFreshGPS(milerUserID) {
			continue
		}

		activeCount := openStopsToday(milerUserID)

		if activeCount >= maxActiveBookings() {
			continue
		}

		onTimeRate, completedToday := fetchMilerStats(milerUserID)
		hubID, hubLoad, hubCapacity := fetchHubData(profile.Hubid)

		candidates = append(candidates, &milerCandidate{
			profile:        profile,
			distanceKm:     loc.Dist,
			activeBookings: activeCount,
			sessionStops:   sessionStops(db.DB, milerUserID),
		})
		aiCandidates = append(aiCandidates, aiCandidate{
			MilerID:        milerUserID,
			DistanceKm:     loc.Dist,
			Rating:         profile.Rating,
			ActiveBookings: activeCount,
			OnTimeRate30d:  onTimeRate,
			CompletedToday: completedToday,
			HubID:          hubID,
			HubLoad:        hubLoad,
			HubCapacity:    hubCapacity,
		})
	}

	return candidates, aiCandidates
}

// openStopsToday is what counts towards the per-rider cap: the rider's open
// assignments (Assigned / Accepted) made since midnight India time, on orders
// that are not cancelled. Before, every open record of any age counted, and
// nothing closed a record when ops cancelled its order — so a rider with a
// pile of old or cancelled orders sat at the cap for good and every new order
// stayed Pending.
func openStopsToday(milerUserID int) int64 {
	return openStopsTodayIn(db.DB, milerUserID)
}

// openStopsTodayIn is openStopsToday on a given handle, so the commit can
// recount inside its own transaction.
func openStopsTodayIn(h *gorm.DB, milerUserID int) int64 {
	var n int64
	h.Table("bookingassignments AS ba").
		Joins("JOIN pickupbookings pb ON pb.bookingid = ba.bookingid").
		Where("ba.mileruserid = ? AND ba.assignmentstatus IN ? AND ba.assignedat >= ? AND pb.status <> ?",
			milerUserID,
			[]string{constants.AssignmentAssigned, constants.AssignmentAccepted},
			startOfISTDay(time.Now()),
			constants.BookingCancelled).
		Count(&n)
	return n
}

// milerHasFreshGPS reports whether the rider's last position is recent enough
// (ASSIGNMENT_MAX_GPS_AGE_MINUTES, default 15; 0 turns the check off).
func milerHasFreshGPS(milerUserID int) bool {
	maxAge := maxGPSAgeMinutes()
	if maxAge == 0 {
		return true
	}
	var n int64
	db.DB.Model(&models.MilerProfile{}).
		Where("userid = ? AND lastlocationupdatedat >= NOW() - make_interval(mins => ?)", milerUserID, maxAge).
		Count(&n)
	return n > 0
}

// sessionStops counts the orders given to the rider since their current duty
// session started (the latest milerdutylogs row with no logout), excluding
// ones that fell through (cancelled, rejected, reassigned). 0 when the rider
// has no open session.
func sessionStops(h *gorm.DB, milerUserID int) int64 {
	var n int64
	h.Table("bookingassignments").
		Where("mileruserid = ? AND assignmentstatus NOT IN ? AND assignedat >= "+
			"(SELECT MAX(loginat) FROM milerdutylogs WHERE userid = ? AND logoutat IS NULL)",
			milerUserID,
			[]string{constants.AssignmentCancelled, constants.AssignmentRejected, constants.AssignmentReassigned},
			milerUserID).
		Count(&n)
	return n
}

// errNoRiderCapacity: by the time the commit ran, every candidate had reached
// the per-rider ceiling. The attempt is retried later like "no rider found".
var errNoRiderCapacity = errors.New("every candidate rider is at the order limit")

// decisionLockKey serialises the final rider choice across attempts and
// replicas (a transaction-scoped Postgres advisory lock). Held only for a
// handful of quick counts and the commit, never across the AI call.
const decisionLockKey int64 = 7_270_001

// finalizeChoice re-checks the choice inside the commit transaction. Attempts
// for different orders run in parallel — a bulk upload of 50 starts 50 — and
// each picked "the least-loaded rider" from counts read before the others
// committed, so they would all pile onto the same one. Under the lock the
// loads are recounted: the chosen rider is kept if still among the least
// loaded and under the ceiling, otherwise the best of the least loaded is
// taken instead (same order as betterChoice).
func finalizeChoice(tx *gorm.DB, chosen *milerCandidate, pool []*milerCandidate) (*milerCandidate, error) {
	if len(pool) == 0 {
		pool = []*milerCandidate{chosen}
	}
	if err := tx.Exec("SELECT pg_advisory_xact_lock(?)", decisionLockKey).Error; err != nil {
		return nil, fmt.Errorf("decision lock: %w", err)
	}
	ceiling := maxActiveBookings()
	var open []*milerCandidate
	for _, c := range pool {
		c.activeBookings = openStopsTodayIn(tx, c.profile.Userid)
		c.sessionStops = sessionStops(tx, c.profile.Userid)
		if c.activeBookings < ceiling {
			open = append(open, c)
		}
	}
	if len(open) == 0 {
		return nil, errNoRiderCapacity
	}
	least, _ := leastLoaded(open, nil)
	for _, c := range least {
		if c.profile.Userid == chosen.profile.Userid {
			return c, nil // the original choice still holds
		}
	}
	return pickBestFromCandidates(least), nil
}

// startOfISTDay is midnight today in India, the boundary for "today's"
// open stops. An absolute instant, so it compares correctly with assignedat
// whether the column stores IST wall-clock or a real timestamp.
func startOfISTDay(now time.Time) time.Time {
	n := now.In(utils.ISTLocation())
	return time.Date(n.Year(), n.Month(), n.Day(), 0, 0, 0, 0, utils.ISTLocation())
}

// ─── Per-miler stats ─────────────────────────────────────────────────────────

func fetchMilerStats(milerUserID int) (onTimeRate float64, completedToday int64) {
	type statsRow struct {
		Total  int64
		OnTime int64
	}
	var row statsRow
	db.DB.Raw(`
		SELECT
			COUNT(*) AS total,
			COUNT(CASE WHEN assignmentstatus = 'Completed_OnTime' THEN 1 END) AS on_time
		FROM bookingassignments
		WHERE mileruserid = ?
		  AND createdat > NOW() - INTERVAL '30 days'
	`, milerUserID).Scan(&row)

	if row.Total == 0 {
		onTimeRate = 0.85 // neutral assumption for milers with no recent history
	} else {
		onTimeRate = float64(row.OnTime) / float64(row.Total)
	}

	today := time.Now().Truncate(24 * time.Hour)
	db.DB.Model(&models.BookingAssignment{}).
		Where("mileruserid = ? AND createdat >= ? AND assignmentstatus = ?",
			milerUserID, today, constants.AssignmentCompleted).
		Count(&completedToday)

	return
}

// ─── Hub stats ───────────────────────────────────────────────────────────────

// fetchHubData returns the hub_load (active assignments across the hub) and
// hub_capacity for the given hub. Defaults to 50 for capacity when the column
// doesn't exist or the miler has no assigned hub.
func fetchHubData(hubID *int) (resolvedHubID int, hubLoad int64, hubCapacity int) {
	hubCapacity = 50 // safe default — also used if hubs.capacity column is absent
	if hubID == nil {
		return 0, 0, hubCapacity
	}
	resolvedHubID = *hubID

	// hub_load: total active assignments whose miler belongs to this hub
	db.DB.Raw(`
		SELECT COUNT(*) FROM bookingassignments ba
		JOIN milerprofiles mp ON ba.mileruserid = mp.userid
		WHERE mp.hubid = ?
		  AND ba.assignmentstatus IN ('Assigned', 'Accepted', 'Pickup_Scheduled')
	`, resolvedHubID).Scan(&hubLoad)

	// hub_capacity: gracefully handle column not existing yet
	var cap int
	if err := db.DB.Raw(`SELECT capacity FROM hubs WHERE hubid = ?`, resolvedHubID).Scan(&cap).Error; err == nil && cap > 0 {
		hubCapacity = cap
	}

	return
}

// ─── Fallback scorer ─────────────────────────────────────────────────────────

// pickBestFromCandidates applies the original formula to an already-collected
// eligible slice: score = distance_km*1.0 + active_bookings*2.0 - rating*0.5
func pickBestFromCandidates(candidates []*milerCandidate) *milerCandidate {
	var best *milerCandidate
	for _, c := range candidates {
		if best == nil || betterChoice(c, best) {
			best = c
		}
	}
	return best
}

// betterChoice is the balancing order: fewest orders in hand, then fewest
// orders this duty session, then nearest to the pickup, then best rated.
// It replaced a weighted score in which distance dominated, so riders near
// the pickups took most of the orders.
func betterChoice(a, b *milerCandidate) bool {
	if a.activeBookings != b.activeBookings {
		return a.activeBookings < b.activeBookings
	}
	if a.sessionStops != b.sessionStops {
		return a.sessionStops < b.sessionStops
	}
	if a.distanceKm != b.distanceKm {
		return a.distanceKm < b.distanceKm
	}
	return a.profile.Rating > b.profile.Rating
}

// ─── AI layer HTTP call ──────────────────────────────────────────────────────

func callDecisionEngine(booking *models.PickupBooking, candidates []aiCandidate) (aiDecisionResponse, error) {
	baseURL := os.Getenv("AI_LAYER_BASE_URL")
	if baseURL == "" {
		baseURL = "https://routemate.workolik.com"
	}

	now := time.Now()
	hour := now.Hour()
	isPeak := (hour >= 7 && hour <= 10) || (hour >= 17 && hour <= 21)

	reqBody := aiDecisionRequest{
		Booking: aiBookingInfo{
			PickupLat:    booking.Pickuplatitude,
			PickupLon:    booking.Pickuplongitude,
			DeliveryLat:  booking.Deliverylatitude,
			DeliveryLon:  booking.Deliverylongitude,
			ItemCategory: b2cFirstParcelCategory(booking),
			Weight:       b2cChargeableWeight(booking),
			ServiceType:  b2cFirstServiceType(booking),
		},
		Candidates: candidates,
		Context: aiRequestCtx{
			Hour:   hour,
			IsPeak: isPeak,
			Zone:   booking.Deliverycity,
		},
	}

	body, err := json.Marshal(reqBody)
	if err != nil {
		return aiDecisionResponse{}, fmt.Errorf("marshal AI request: %w", err)
	}

	utils.Info("Assignment: AI layer outgoing payload",
		"booking_id", booking.Bookingid,
		"payload", string(body),
	)

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	req, err := http.NewRequestWithContext(ctx, http.MethodPost,
		baseURL+"/api/v1/doormile/decide-assignment",
		bytes.NewReader(body))
	if err != nil {
		return aiDecisionResponse{}, fmt.Errorf("build AI request: %w", err)
	}
	req.Header.Set("Content-Type", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return aiDecisionResponse{}, fmt.Errorf("AI layer call: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return aiDecisionResponse{}, fmt.Errorf("AI layer HTTP %d", resp.StatusCode)
	}

	var decision aiDecisionResponse
	if err := json.NewDecoder(resp.Body).Decode(&decision); err != nil {
		return aiDecisionResponse{}, fmt.Errorf("decode AI response: %w", err)
	}

	return decision, nil
}