package main

import (
	"context"
	"encoding/json"
	"fmt"
	"sort"
	"strings"
	"sync"
)

var includedMethods = map[string]bool{
	"walk":       true,
	"surf":       true,
	"old-rod":    true,
	"good-rod":   true,
	"super-rod":  true,
	"rock-smash": true,
	"headbutt":   true,
}

// processVersion processes all locations for a game version and returns the route list.
func processVersion(
	ctx context.Context,
	client *Client,
	versionName string,
	vgInfo VersionGroupInfo,
	vgKey string,
	routeOrder map[string][]string,
	specialEnc *SpecialEncountersFile,
	pokeIDCollector *PokeIDCollector,
) ([]RouteOutput, error) {
	fmt.Printf("\n--- Processing %s ---\n", versionName)

	// Fetch region
	regionData, err := client.Get(ctx, fmt.Sprintf("region/%d", vgInfo.RegionID))
	if err != nil {
		return nil, fmt.Errorf("fetching region %d: %w", vgInfo.RegionID, err)
	}
	var region RegionResp
	if err := json.Unmarshal(regionData, &region); err != nil {
		return nil, fmt.Errorf("parsing region %d: %w", vgInfo.RegionID, err)
	}

	locationRefs := make([]NamedRef, len(region.Locations))
	copy(locationRefs, region.Locations)

	// Include extra regions
	for _, extraRegionID := range vgInfo.ExtraRegions {
		extraData, err := client.Get(ctx, fmt.Sprintf("region/%d", extraRegionID))
		if err != nil {
			return nil, fmt.Errorf("fetching extra region %d: %w", extraRegionID, err)
		}
		var extraRegion RegionResp
		if err := json.Unmarshal(extraData, &extraRegion); err != nil {
			return nil, fmt.Errorf("parsing extra region %d: %w", extraRegionID, err)
		}
		locationRefs = append(locationRefs, extraRegion.Locations...)
	}

	fmt.Printf("  Found %d locations\n", len(locationRefs))

	// Fetch all locations concurrently
	type locationResult struct {
		locName string
		locID   int
		areas   []NamedRef
	}

	locResults := make([]locationResult, len(locationRefs))
	var wg sync.WaitGroup
	errs := make([]error, len(locationRefs))

	for i, locRef := range locationRefs {
		wg.Add(1)
		go func(i int, locRef NamedRef) {
			defer wg.Done()
			locData, err := client.Get(ctx, fmt.Sprintf("location/%d", locRef.ID()))
			if err != nil {
				errs[i] = err
				return
			}
			var loc LocationResp
			if err := json.Unmarshal(locData, &loc); err != nil {
				errs[i] = err
				return
			}
			locResults[i] = locationResult{
				locName: locRef.Name,
				locID:   locRef.ID(),
				areas:   loc.Areas,
			}
		}(i, locRef)
	}
	wg.Wait()

	for _, err := range errs {
		if err != nil {
			return nil, err
		}
	}

	// Fetch all area encounters concurrently
	type areaWork struct {
		locIdx    int
		areaRef   NamedRef
		locName   string
		areaCount int // total areas for this location
	}

	var areaJobs []areaWork
	for i, lr := range locResults {
		for _, areaRef := range lr.areas {
			areaJobs = append(areaJobs, areaWork{
				locIdx:    i,
				areaRef:   areaRef,
				locName:   lr.locName,
				areaCount: len(lr.areas),
			})
		}
	}

	type areaResult struct {
		locIdx     int
		areaSuffix string
		areaCount  int
		encounters []EncounterOutput
	}

	areaResults := make([]areaResult, len(areaJobs))
	areaErrs := make([]error, len(areaJobs))

	for i, job := range areaJobs {
		wg.Add(1)
		go func(i int, job areaWork) {
			defer wg.Done()
			encs, err := getEncountersForArea(ctx, client, job.areaRef.ID(), versionName)
			if err != nil {
				areaErrs[i] = err
				return
			}
			areaSuffix := CleanAreaName(job.areaRef.Name, job.locName)
			areaResults[i] = areaResult{
				locIdx:     job.locIdx,
				areaSuffix: areaSuffix,
				areaCount:  job.areaCount,
				encounters: encs,
			}
		}(i, job)
	}
	wg.Wait()

	for _, err := range areaErrs {
		if err != nil {
			return nil, err
		}
	}

	// Group area results by location
	type locAreaData struct {
		allEncounters []EncounterOutput
		areaSpecific  map[string][]EncounterOutput
	}
	locAreas := make(map[int]*locAreaData)
	for _, ar := range areaResults {
		ld, ok := locAreas[ar.locIdx]
		if !ok {
			ld = &locAreaData{areaSpecific: make(map[string][]EncounterOutput)}
			locAreas[ar.locIdx] = ld
		}
		if len(ar.encounters) == 0 {
			continue
		}
		if ar.areaSuffix != "" && ar.areaCount > 1 {
			ld.areaSpecific[ar.areaSuffix] = append(ld.areaSpecific[ar.areaSuffix], ar.encounters...)
		} else {
			ld.allEncounters = append(ld.allEncounters, ar.encounters...)
		}
	}

	// Build routes
	var routes []RouteOutput
	for i, lr := range locResults {
		if len(lr.areas) == 0 {
			continue
		}
		displayName := CleanLocationName(lr.locName)
		ld, ok := locAreas[i]
		if !ok {
			continue
		}

		// Multiple area-specific encounters -> parent with children
		if len(ld.areaSpecific) > 1 {
			var childRoutes []RouteOutput
			for areaSuffix, areaEncs := range ld.areaSpecific {
				aggregated := aggregateEncounters(areaEncs)
				if len(aggregated) > 0 {
					routeName := fmt.Sprintf("%s (%s)", displayName, areaSuffix)
					for _, enc := range aggregated {
						pokeIDCollector.Add(enc.PokeAPIID)
					}
					childRoutes = append(childRoutes, RouteOutput{
						Name:       routeName,
						Order:      0,
						Encounters: aggregated,
					})
				}
			}
			if len(childRoutes) > 0 {
				routes = append(routes, RouteOutput{
					Name:       displayName,
					Order:      0,
					Encounters: []EncounterOutput{},
					Children:   childRoutes,
				})
			}
		} else if len(ld.areaSpecific) == 1 {
			// Single area-specific -> no parent/child
			for areaSuffix, areaEncs := range ld.areaSpecific {
				aggregated := aggregateEncounters(areaEncs)
				if len(aggregated) > 0 {
					routeName := fmt.Sprintf("%s (%s)", displayName, areaSuffix)
					for _, enc := range aggregated {
						pokeIDCollector.Add(enc.PokeAPIID)
					}
					routes = append(routes, RouteOutput{
						Name:       routeName,
						Order:      0,
						Encounters: aggregated,
					})
				}
			}
		}

		// Non-area-specific encounters
		if len(ld.allEncounters) > 0 {
			aggregated := aggregateEncounters(ld.allEncounters)
			if len(aggregated) > 0 {
				for _, enc := range aggregated {
					pokeIDCollector.Add(enc.PokeAPIID)
				}
				routes = append(routes, RouteOutput{
					Name:       displayName,
					Order:      0,
					Encounters: aggregated,
				})
			}
		}
	}

	// Merge special encounters
	specialData := getSpecialEncounters(specialEnc, vgKey)
	if specialData != nil {
		routes = mergeSpecialEncounters(routes, specialData, pokeIDCollector)
	}

	// Sort by game progression
	routes = sortRoutesByProgression(routes, vgKey, routeOrder)

	// Assign sequential order values
	order := 1
	for i := range routes {
		routes[i].Order = order
		order++
		for j := range routes[i].Children {
			routes[i].Children[j].Order = order
			order++
		}
	}

	// Stats
	totalRoutes := 0
	totalEnc := 0
	for _, r := range routes {
		totalRoutes += 1 + len(r.Children)
		totalEnc += len(r.Encounters)
		for _, c := range r.Children {
			totalEnc += len(c.Encounters)
		}
	}
	fmt.Printf("  Routes with encounters: %d\n", totalRoutes)
	fmt.Printf("  Total encounter entries: %d\n", totalEnc)

	return routes, nil
}

// getEncountersForArea fetches encounter data for a location area, filtered by version.
func getEncountersForArea(
	ctx context.Context,
	client *Client,
	areaID int,
	versionName string,
) ([]EncounterOutput, error) {
	data, err := client.Get(ctx, fmt.Sprintf("location-area/%d", areaID))
	if err != nil {
		return nil, err
	}
	var area LocationAreaResp
	if err := json.Unmarshal(data, &area); err != nil {
		return nil, err
	}

	var encounters []EncounterOutput
	for _, pe := range area.PokemonEncounters {
		dexNum := pe.Pokemon.ID()
		pokemonName := pe.Pokemon.Name

		for _, vd := range pe.VersionDetails {
			if vd.Version.Name != versionName {
				continue
			}
			for _, enc := range vd.EncounterDetails {
				method := enc.Method.Name
				if !includedMethods[method] {
					continue
				}
				encounters = append(encounters, EncounterOutput{
					PokemonName:   pokemonName,
					PokeAPIID:     dexNum,
					Method:        method,
					EncounterRate: enc.Chance,
					MinLevel:      enc.MinLevel,
					MaxLevel:      enc.MaxLevel,
				})
			}
		}
	}
	return encounters, nil
}

// aggregateEncounters groups encounters by (pokeapi_id, method) and sums rates.
func aggregateEncounters(raw []EncounterOutput) []EncounterOutput {
	type key struct {
		id     int
		method string
	}
	agg := make(map[key]*EncounterOutput)
	var order []key // preserve insertion order

	for _, enc := range raw {
		k := key{enc.PokeAPIID, enc.Method}
		if existing, ok := agg[k]; ok {
			existing.EncounterRate += enc.EncounterRate
			if enc.MinLevel < existing.MinLevel {
				existing.MinLevel = enc.MinLevel
			}
			if enc.MaxLevel > existing.MaxLevel {
				existing.MaxLevel = enc.MaxLevel
			}
		} else {
			e := enc // copy
			agg[k] = &e
			order = append(order, k)
		}
	}

	result := make([]EncounterOutput, 0, len(agg))
	for _, k := range order {
		e := agg[k]
		if e.EncounterRate > 100 {
			e.EncounterRate = 100
		}
		result = append(result, *e)
	}

	sort.Slice(result, func(i, j int) bool {
		if result[i].EncounterRate != result[j].EncounterRate {
			return result[i].EncounterRate > result[j].EncounterRate
		}
		return result[i].PokemonName < result[j].PokemonName
	})

	return result
}

// mergeSpecialEncounters merges special encounters into existing routes or creates new ones.
func mergeSpecialEncounters(
	routes []RouteOutput,
	specialData map[string][]EncounterOutput,
	pokeIDCollector *PokeIDCollector,
) []RouteOutput {
	// Build lookup: route name -> route pointer (including children)
	routeMap := make(map[string]*RouteOutput)
	for i := range routes {
		routeMap[routes[i].Name] = &routes[i]
		for j := range routes[i].Children {
			routeMap[routes[i].Children[j].Name] = &routes[i].Children[j]
		}
	}

	for locationName, encounters := range specialData {
		for _, enc := range encounters {
			pokeIDCollector.Add(enc.PokeAPIID)
		}
		if route, ok := routeMap[locationName]; ok {
			route.Encounters = append(route.Encounters, encounters...)
		} else {
			newRoute := RouteOutput{
				Name:       locationName,
				Order:      0,
				Encounters: encounters,
			}
			routes = append(routes, newRoute)
			routeMap[locationName] = &routes[len(routes)-1]
		}
	}

	return routes
}

// sortRoutesByProgression sorts routes by game progression order.
func sortRoutesByProgression(routes []RouteOutput, vgKey string, routeOrder map[string][]string) []RouteOutput {
	orderList, ok := routeOrder[vgKey]
	if !ok {
		return routes
	}

	sort.SliceStable(routes, func(i, j int) bool {
		iKey := routeSortKey(routes[i].Name, orderList)
		jKey := routeSortKey(routes[j].Name, orderList)
		if iKey.pos != jKey.pos {
			return iKey.pos < jKey.pos
		}
		return iKey.name < jKey.name
	})

	return routes
}

type sortKey struct {
	pos  int
	name string
}

func routeSortKey(name string, orderList []string) sortKey {
	for i, orderedName := range orderList {
		if name == orderedName || strings.HasPrefix(name, orderedName+" (") {
			return sortKey{i, name}
		}
	}
	return sortKey{len(orderList), name}
}

// PokeIDCollector is a thread-safe collector for PokeAPI IDs encountered during processing.
type PokeIDCollector struct {
	mu  sync.Mutex
	ids map[int]bool
}

func NewPokeIDCollector() *PokeIDCollector {
	return &PokeIDCollector{ids: make(map[int]bool)}
}

func (c *PokeIDCollector) Add(id int) {
	c.mu.Lock()
	c.ids[id] = true
	c.mu.Unlock()
}

func (c *PokeIDCollector) IDs() map[int]bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	result := make(map[int]bool, len(c.ids))
	for k, v := range c.ids {
		result[k] = v
	}
	return result
}