package main

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

// fetchAllSpecies fetches all pokemon-species objects.
// Returns a map of species_id -> SpeciesResp.
func fetchAllSpecies(ctx context.Context, client *Client) (map[int]*SpeciesResp, error) {
	// Fetch the species list
	listData, err := client.Get(ctx, "pokemon-species?limit=10000")
	if err != nil {
		return nil, fmt.Errorf("fetching species list: %w", err)
	}
	var listing SpeciesListResp
	if err := json.Unmarshal(listData, &listing); err != nil {
		return nil, fmt.Errorf("parsing species list: %w", err)
	}

	// Filter to IDs < 10000 and sort
	var speciesIDs []int
	for _, entry := range listing.Results {
		id := entry.ID()
		if id < 10000 {
			speciesIDs = append(speciesIDs, id)
		}
	}
	sort.Ints(speciesIDs)

	fmt.Printf("\n--- Fetching %d species data ---\n", len(speciesIDs))

	speciesData := make(map[int]*SpeciesResp, len(speciesIDs))
	var mu sync.Mutex
	var wg sync.WaitGroup
	errs := make([]error, len(speciesIDs))

	for i, sid := range speciesIDs {
		wg.Add(1)
		go func(i, sid int) {
			defer wg.Done()
			data, err := client.Get(ctx, fmt.Sprintf("pokemon-species/%d", sid))
			if err != nil {
				errs[i] = err
				return
			}
			var species SpeciesResp
			if err := json.Unmarshal(data, &species); err != nil {
				errs[i] = fmt.Errorf("parsing species %d: %w", sid, err)
				return
			}
			mu.Lock()
			speciesData[sid] = &species
			mu.Unlock()
		}(i, sid)

		// Print progress every 200
		if (i+1)%200 == 0 || i+1 == len(speciesIDs) {
			// Progress will be approximate due to concurrency
		}
	}
	wg.Wait()

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

	fmt.Printf("  Fetched %d/%d species\n", len(speciesData), len(speciesIDs))
	return speciesData, nil
}

// fetchAllPokemon fetches all Pokemon (base + forms) and returns sorted output.
func fetchAllPokemon(
	ctx context.Context,
	client *Client,
	speciesData map[int]*SpeciesResp,
	allPokeAPIIDs map[int]bool,
) ([]PokemonOutput, error) {
	// Collect base species IDs and form IDs from species varieties
	var baseIDs []int
	var formIDs []int
	formIDSet := make(map[int]bool)

	for _, species := range speciesData {
		for _, variety := range species.Varieties {
			pid := variety.Pokemon.ID()
			if variety.IsDefault {
				baseIDs = append(baseIDs, pid)
			} else {
				formIDs = append(formIDs, pid)
				formIDSet[pid] = true
			}
		}
	}

	// Also include form IDs from encounter data not in varieties
	for id := range allPokeAPIIDs {
		if id >= 10000 && !formIDSet[id] {
			formIDs = append(formIDs, id)
		}
	}

	sort.Ints(baseIDs)
	sort.Ints(formIDs)

	fmt.Printf("\n--- Fetching %d base Pokemon + %d forms ---\n", len(baseIDs), len(formIDs))

	// Fetch base Pokemon concurrently
	type pokemonResult struct {
		output PokemonOutput
		isForm bool
	}

	allIDs := make([]int, 0, len(baseIDs)+len(formIDs))
	isFormFlag := make([]bool, 0, len(baseIDs)+len(formIDs))
	for _, id := range baseIDs {
		allIDs = append(allIDs, id)
		isFormFlag = append(isFormFlag, false)
	}
	for _, id := range formIDs {
		allIDs = append(allIDs, id)
		isFormFlag = append(isFormFlag, true)
	}

	results := make([]pokemonResult, len(allIDs))
	var wg sync.WaitGroup
	errs := make([]error, len(allIDs))

	for i, pid := range allIDs {
		wg.Add(1)
		go func(i, pid int, isForm bool) {
			defer wg.Done()
			data, err := client.Get(ctx, fmt.Sprintf("pokemon/%d", pid))
			if err != nil {
				errs[i] = err
				return
			}
			var poke PokemonResp
			if err := json.Unmarshal(data, &poke); err != nil {
				errs[i] = fmt.Errorf("parsing pokemon %d: %w", pid, err)
				return
			}

			var types []string
			for _, t := range poke.Types {
				types = append(types, t.Type.Name)
			}

			var name string
			var nationalDex int
			if isForm {
				name = FormatFormName(poke.Name, poke.Species.Name)
				nationalDex = poke.Species.ID()
			} else {
				name = toTitleCase(poke.Name)
				nationalDex = pid
			}

			results[i] = pokemonResult{
				output: PokemonOutput{
					PokeAPIID:   pid,
					NationalDex: nationalDex,
					Name:        name,
					Types:       types,
					SpriteURL:   fmt.Sprintf("https://raw.githubusercontent.com/PokeAPI/sprites/master/sprites/pokemon/%d.png", pid),
				},
				isForm: isForm,
			}
		}(i, pid, isFormFlag[i])
	}
	wg.Wait()

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

	pokemonList := make([]PokemonOutput, 0, len(results))
	for _, r := range results {
		pokemonList = append(pokemonList, r.output)
	}

	sort.Slice(pokemonList, func(i, j int) bool {
		if pokemonList[i].NationalDex != pokemonList[j].NationalDex {
			return pokemonList[i].NationalDex < pokemonList[j].NationalDex
		}
		return pokemonList[i].PokeAPIID < pokemonList[j].PokeAPIID
	})

	fmt.Printf("  Fetched %d base Pokemon\n", len(baseIDs))
	fmt.Printf("  Fetched %d forms\n", len(formIDs))

	return pokemonList, nil
}
