package main

import (
	"container/heap"
	"fmt"
	"log"
	"time"

	"git.tornberg.me/go-gtfs/pkg/reader"
	"git.tornberg.me/go-gtfs/pkg/types"
)

// TripPlanner handles preprocessed transit data for efficient routing
type TripPlanner struct {
	*reader.TripData
}

const (
	transferPenalty     = 90 * time.Minute
	maxTransfers        = 4
	maxWaitBetweenTrips = 1 * time.Hour
	//trajectoryAngleTolerance = 220.0
	maxTravelDuration = 12 * time.Hour
	maxInitialWait    = 90 * time.Minute
	searchWindow      = 4 * time.Hour
)

// NewTripPlanner creates a new trip planner instance
func NewTripPlanner(data *reader.TripData) *TripPlanner {
	return &TripPlanner{
		TripData: data,
		//graph:    make(map[string][]Edge),
	}
}

// Preprocess builds the routing graph and precomputes routes
func (tp *TripPlanner) Preprocess() error {

	// if hosjo, ok := tp.Stops["740025287"]; ok {
	// 	trips := hosjo.GetTripsAfter(time.Now())
	// 	for trip := range trips {
	// 		log.Printf("Trip %s (%s):", trip.TripShortName, trip.TripHeadsign)
	// 		for stop := range trip.GetDirectPossibleDestinations(hosjo, time.Now()) {
	// 			log.Printf("- Stop %s at %s", stop.Stop.StopName, types.AsTime(stop.DepartureTime))
	// 		}
	// 	}

	// }

	// // Build graph with trip edges
	// for tripID, trip := range tp.Trips {
	// 	sts := trip.Stops
	// 	sort.Slice(sts, func(i, j int) bool {
	// 		return sts[i].StopSequence < sts[j].StopSequence
	// 	})
	// 	for i := 0; i < len(sts)-1; i++ {
	// 		from := sts[i].StopId
	// 		to := sts[i+1].StopId
	// 		departure := sts[i].DepartureTime
	// 		arrival := sts[i+1].DepartureTime
	// 		timeDiff := arrival - departure
	// 		if timeDiff > 0 {
	// 			tp.graph[from] = append(tp.graph[from], Edge{To: to, TripID: tripID, Time: timeDiff, DepartureTime: departure})
	// 		}
	// 	}
	// }

	// // Add transfer edges
	// for _, tr := range tp.transfers {
	// 	if tr.TransferType == 2 { // minimum transfer time
	// 		tp.graph[tr.FromStopId] = append(tp.graph[tr.FromStopId], Edge{
	// 			To:            tr.ToStopId,
	// 			TripID:        "transfer",
	// 			Time:          float64(tr.MinTransferTime),
	// 			DepartureTime: 0,
	// 		})
	// 	}
	// }

	return nil
}

type searchNode struct {
	stopTime  *types.StopTime
	index     int
	g         float64
	f         float64
	transfers int
	parent    *searchNode
}

func nodeKey(st *types.StopTime) string {
	return fmt.Sprintf("%s:%d", st.TripId, st.StopSequence)
}

func heuristicSeconds(from, to *types.Stop) float64 {
	if from == nil || to == nil {
		return 0
	}
	const avgSpeedKmh = 60.0
	distanceKm := from.HaversineDistance(to)
	return (distanceKm / avgSpeedKmh) * 3600
}

func tripStopIndex(trip *types.Trip, st *types.StopTime) int {
	if trip == nil {
		return -1
	}
	for i, candidate := range trip.Stops {
		if candidate == st {
			return i
		}
	}
	return -1
}

// FindRoutes finds the best routes (up to num) between two stops starting at the given time
func (tp *TripPlanner) FindRoute(from, to string, when time.Time) ([]*Route, error) {
	fromStop := tp.GetStop(from)
	toStop := tp.GetStop(to)
	if fromStop == nil || toStop == nil {
		return nil, fmt.Errorf("invalid from or to stop")
	}

	exclude := make(map[string]struct{})
	results := make([]*Route, 0, 3)
	timeCursor := when
	const maxAttempts = 5
	var lastErr error

	for attempts := 0; attempts < maxAttempts && len(results) < 3; attempts++ {
		limit := 3 - len(results)
		routes, err := tp.searchRoutesAStar(fromStop, toStop, timeCursor, exclude, limit)
		if err == nil {
			results = append(results, routes...)
		} else {
			lastErr = err
		}
		timeCursor = timeCursor.Add(10 * time.Minute)
	}

	if len(results) == 0 {
		if lastErr != nil {
			return nil, lastErr
		}
		return nil, fmt.Errorf("no route found")
	}

	return results, nil
}

func (tp *TripPlanner) searchRoutesAStar(fromStop, toStop *types.Stop, when time.Time, exclude map[string]struct{}, limit int) ([]*Route, error) {
	if limit <= 0 {
		return nil, nil
	}
	open := priorityQueue{}
	heap.Init(&open)
	bestCost := make(map[string]float64)
	startSeconds := types.AsSecondsAfterMidnight(when)
	startSecondsFloat := float64(startSeconds)
	initCount := 0
	initialWaitLimit := maxInitialWait.Seconds()
	searchWindowLimit := searchWindow.Seconds()

	for tripWithDeparture := range fromStop.GetTripsAfter(when) {
		trip := tripWithDeparture.Trip
		startStopTime, ok := trip.Has(fromStop)
		if !ok {
			continue
		}
		idx := tripStopIndex(trip, startStopTime)
		if idx < 0 {
			continue
		}
		departureDelta := float64(startStopTime.DepartureTime) - startSecondsFloat
		if searchWindow > 0 && departureDelta > searchWindowLimit {
			continue
		}
		wait := departureDelta
		if wait < 0 {
			wait = 0
		}
		if maxInitialWait > 0 && wait > initialWaitLimit {
			continue
		}
		node := &searchNode{
			stopTime:  startStopTime,
			index:     idx,
			g:         wait,
			transfers: 0,
		}
		node.f = node.g + heuristicSeconds(startStopTime.Stop, toStop)
		heap.Push(&open, &pqItem{node: node, priority: node.f})
		bestCost[nodeKey(startStopTime)] = node.g
		initCount++
	}

	if initCount == 0 {
		return nil, fmt.Errorf("no departures from %s after %s", fromStop.StopName, when.Format(time.RFC3339))
	}

	routes := make([]*Route, 0, limit)

	for open.Len() > 0 && len(routes) < limit {
		current := heap.Pop(&open).(*pqItem).node
		if current.stopTime.StopId == toStop.StopId {
			route := buildRouteFromNode(current)
			if route != nil {
				signature := routeSignature(route)
				if signature == "" {
					continue
				}
				if _, seen := exclude[signature]; seen {
					continue
				}
				exclude[signature] = struct{}{}
				routes = append(routes, route)
				if len(routes) >= limit {
					break
				}
			}
			continue
		}
		if current.g > maxTravelDuration.Seconds() {
			continue
		}
		if searchWindow > 0 {
			currentDelta := float64(current.stopTime.DepartureTime) - startSecondsFloat
			if currentDelta > searchWindowLimit {
				continue
			}
		}

		tpp := tp.Trips[current.stopTime.TripId]
		if tpp != nil {
			for i := current.index + 1; i < len(tpp.Stops); i++ {
				next := tpp.Stops[i]
				if next.DropOffType == 1 {
					continue
				}
				travel := next.ArrivalTime - current.stopTime.DepartureTime
				if travel <= 0 {
					continue
				}
				newCost := current.g + float64(travel)
				if newCost > maxTravelDuration.Seconds() {
					continue
				}
				if searchWindow > 0 {
					nextDepartureDelta := float64(next.DepartureTime) - startSecondsFloat
					if nextDepartureDelta > searchWindowLimit {
						continue
					}
				}
				key := nodeKey(next)
				if prev, ok := bestCost[key]; ok && newCost >= prev {
					continue
				}
				nextNode := &searchNode{
					stopTime:  next,
					index:     i,
					g:         newCost,
					transfers: current.transfers,
					parent:    current,
				}
				nextNode.f = newCost + heuristicSeconds(next.Stop, toStop)
				bestCost[key] = newCost
				heap.Push(&open, &pqItem{node: nextNode, priority: nextNode.f})
			}
		}
		if current.transfers >= maxTransfers {
			continue
		}
		for _, otherTrip := range current.stopTime.Stop.Trips {
			if otherTrip.TripId == current.stopTime.TripId {
				continue
			}
			transferStopTime, ok := otherTrip.Has(current.stopTime.Stop)
			if !ok {
				continue
			}
			if transferStopTime.DepartureTime <= current.stopTime.ArrivalTime {
				continue
			}
			wait := transferStopTime.DepartureTime - current.stopTime.ArrivalTime
			if wait > types.SecondsAfterMidnight(maxWaitBetweenTrips.Seconds()) {
				continue
			}
			newCost := current.g + float64(wait) + transferPenalty.Seconds()
			if newCost > maxTravelDuration.Seconds() {
				continue
			}
			if searchWindow > 0 {
				transferDepartureDelta := float64(transferStopTime.DepartureTime) - startSecondsFloat
				if transferDepartureDelta > searchWindowLimit {
					continue
				}
			}
			idx := tripStopIndex(otherTrip, transferStopTime)
			if idx < 0 {
				continue
			}
			key := nodeKey(transferStopTime)
			if prev, ok := bestCost[key]; ok && newCost >= prev {
				continue
			}
			nextNode := &searchNode{
				stopTime:  transferStopTime,
				index:     idx,
				g:         newCost,
				transfers: current.transfers + 1,
				parent:    current,
			}
			nextNode.f = newCost + heuristicSeconds(transferStopTime.Stop, toStop)
			bestCost[key] = newCost
			heap.Push(&open, &pqItem{node: nextNode, priority: nextNode.f})
		}
	}

	if len(routes) == 0 {
		return nil, fmt.Errorf("no route found")
	}

	return routes, nil
}

func buildRouteFromNode(goal *searchNode) *Route {
	if goal == nil {
		return nil
	}
	nodes := make([]*searchNode, 0)
	for current := goal; current != nil; current = current.parent {
		nodes = append(nodes, current)
	}
	for i, j := 0, len(nodes)-1; i < j; i, j = i+1, j-1 {
		nodes[i], nodes[j] = nodes[j], nodes[i]
	}
	if len(nodes) < 2 {
		return nil
	}
	legs := make([]Leg, 0)
	runStart := nodes[0]
	for i := 1; i < len(nodes); i++ {
		prev := nodes[i-1]
		curr := nodes[i]
		if curr.stopTime.TripId != prev.stopTime.TripId {
			legs = append(legs, NewLeg(runStart.stopTime, prev.stopTime))
			runStart = curr
		}
	}
	legs = append(legs, NewLeg(runStart.stopTime, nodes[len(nodes)-1].stopTime))
	return &Route{Legs: legs}
}

func routeSignature(route *Route) string {
	if route == nil {
		return ""
	}
	if len(route.Legs) == 0 {
		return ""
	}
	start := int(route.StartTime())
	duration := route.Duration()
	return fmt.Sprintf("%d:%d", start, duration)
}

func NewLeg(fromStop, toStop *types.StopTime) Leg {
	trip, ok := toStop.Stop.Trips[toStop.TripId]
	if !ok {
		log.Printf("trip %s not found for stop %s", toStop.TripId, toStop.Stop.StopName)
		return Leg{
			From: fromStop,
			To:   toStop,
		}
	}
	agencyName := ""
	if trip.Agency != nil {
		agencyName = trip.Agency.AgencyName
	}
	return Leg{
		From: fromStop,
		To:   toStop,
		Trip: &JSONTrip{
			TripId:        trip.TripId,
			RouteId:       trip.RouteId,
			AgencyName:    agencyName,
			TripHeadsign:  trip.TripHeadsign,
			TripShortName: trip.TripShortName,
		},
	}
}

// // findRoute implements a time-aware Dijkstra algorithm for routing
// func (tp *TripPlanner) findRoute(start, end string, when time.Time) *Route {
// 	csaPlanner := NewCSAPlanner(tp.TripData)
// 	return csaPlanner.FindRoute(start, end, when)
// }

func (tp *TripPlanner) GetRoute(routeId string) *types.Route {
	if routeId == "" {
		return nil
	}
	route, ok := tp.Routes[routeId]
	if !ok {
		return nil
	}
	return route
}

func (tp *TripPlanner) GetAgency(agencyId string) *types.Agency {
	if agencyId == "" {
		return nil
	}
	agency, ok := tp.Agencies[agencyId]
	if !ok {
		return nil
	}
	return agency
}

func (tp *TripPlanner) GetTrip(tripId string) *types.Trip {
	if tripId == "" {
		return nil
	}
	trip, ok := tp.Trips[tripId]
	if !ok {
		return nil
	}
	return trip
}

func (tp *TripPlanner) GetStop(prev string) *types.Stop {
	if prev == "" {
		return nil
	}
	stop, ok := tp.Stops[prev]
	if !ok {
		return nil
	}
	return stop
}