package planet

import (
	"fmt"
	"sort"
	"sync"
	"time"

	"salty/terrain/internal/coast"
	"salty/terrain/internal/field"
	"salty/terrain/internal/fluvial"
	"salty/terrain/internal/manifest"
	"salty/terrain/internal/region"
	"salty/terrain/internal/stats"
	"salty/terrain/internal/template"
	"salty/terrain/internal/thermal"
	"salty/terrain/internal/uplift"
)

// Bake solves the geology of every region and composites the result into one planet.
//
// The order is the whole design: each landmass is solved in a box of its own with water all round it, which
// is exactly the same answer as solving the planet whole because no flow path crosses open water; then the
// land is written back; then the sea floor is laid once, over the finished cylinder.

// Result is a baked planet at geology resolution.
type Result struct {
	In     *Inputs
	Height *field.Field // metres, over the whole planet including the polar pad
	Flow   []float32    // drainage area in m2, for the rivers on the preview
	Sea    []bool       // after the bake: below sea level

	Regions []RegionResult
	Craters []CraterStats
	Elapsed time.Duration

	// Coast is the shelf, the surf and the sediment budget, run once over the whole cylinder after every
	// region is composited. Nil only when there is no planet to run it on.
	Coast *coast.Result

	// Stats is the whole planet's, pooled from the regions. Nil until Write computes it.
	//
	// Pooled rather than averaged, which is the one thing that makes it a planet statistic at all: a region
	// holds a histogram and histograms add, so merging them and taking a quantile of the sum gives exactly
	// what one pass over the whole world would have. See internal/stats.
	Stats *stats.Report

	cancelled bool
}

// RegionResult is what one region's solve cost and produced.
type RegionResult struct {
	ID        int
	Cells     int
	LandCells int
	Seconds   float64
	MinM      float64
	MaxM      float64
	ClipFrac  float64

	// FaultClamped is how many cells the fault pass pushed past the angle of repose and had to bound. Not an
	// error - it is the set saying the throws are large for the class they sit in - but worth seeing in the
	// summary rather than discovering later in a hillshade full of polygonal facets.
	FaultClamped int

	// stats is this region's land statistics, unexported because it is scaffolding: it exists to be merged
	// into the planet's and is not part of what a region result means.
	stats *stats.Accumulator
}

// statsOptions is what a world is judged against, in one place so a region and the planet it belongs to
// cannot disagree about it - two accumulators built on different bounds do not merge.
func statsOptions(m *manifest.Manifest) stats.Options {
	return stats.Options{
		ElevMin: m.ElevationM.Min, ElevMax: m.ElevationM.Max,
		TalusDeg: m.Pipeline.Thermal.TalusDeg, ReliefWindowM: reliefWindowM,
		ChannelM2: channelKm2 * 1e6,
		K:         m.Pipeline.Fluvial.K, M: m.Pipeline.Fluvial.M, N: m.Pipeline.Fluvial.N,
	}
}

// The two constants the statistics are taken at. They are flags on `generate` and fixed here, because a bake
// is compared against other bakes and a window that moved between them would make the relief column
// meaningless: 500 m is the usual choice, and 1 km2 is the incoming spec's channel definition.
const (
	reliefWindowM = 500.0
	channelKm2    = 1.0
)

// BakeOptions steer a run without editing the manifest.
type BakeOptions struct {
	Only  []int // region ids; empty means all
	Steps int   // override the fluvial step count
	Jobs  int   // regions solved at once; 0 is the default
	Log   func(string, ...any)

	// OnRegion is called each time a region's land has been written back, with the planet as it stands.
	//
	// It runs holding the composite lock, so Height and Flow can be read without racing the workers still
	// running - and every one of them is stopped while it does, so it has to be quick and it must not keep a
	// reference to either past the call. It exists so a caller can *watch* a bake: two hours is a long time
	// to find out at the end that the numbers were wrong, and the world filling in one landmass at a time
	// answers that at the first one.
	OnRegion func(*Result, RegionResult)

	// Cancel abandons the run when it is closed. Regions in flight stop at the end of their current step and
	// their land is composited in whatever state the solve had reached, so a cancelled Result is for looking
	// at and never for writing out as a bake.
	Cancel <-chan struct{}
}

// Cancelled reports whether a run was abandoned before it finished.
func (r *Result) Cancelled() bool { return r.cancelled }

// Bake runs the geology solve over a prepared planet.
func Bake(in *Inputs, opt BakeOptions) (*Result, error) {
	log := opt.Log
	if log == nil {
		log = func(string, ...any) {}
	}
	m := in.M
	p := in.P
	n := p.W * p.H
	started := time.Now()

	params := solveParams(m)
	if opt.Steps > 0 {
		params.Steps = opt.Steps
	}
	log("fluvial %d steps of %.0f yr (%.2f Myr), K %.1e, m %.2f, n %.2f, fill every %d",
		params.Steps, params.DtYr, float64(params.Steps)*params.DtYr/1e6,
		params.K, params.M, params.N, params.FillEvery)

	res := &Result{In: in, Height: field.New(p.W, p.H, p.CellM), Flow: make([]float32, n)}
	for i := range res.Height.Data {
		res.Height.Data[i] = float32(m.SeaLevelM)
	}

	rates := in.Legend.Rates()
	ks := in.Legend.Erodibilities()
	plainM, plainFloor := in.Legend.CoastalPlains()
	massifFloor, massifFraction := in.Legend.Massifs()
	wanted := wantedRegions(in.Part.Regions, opt.Only)

	// Regions are solved a few at a time. They are independent - each writes only the land it owns, and no
	// flow path crosses the water between them - so this changes nothing about the result; what it buys is
	// the parts of the solve that are sequential *within* a region. Terrain.md's profile says most of the
	// runtime is the stack walk and the flood's cursor, neither of which parallelises inside one grid, so
	// overlapping regions is where the cores actually go.
	//
	// Results land in indexed slots and are read back in region order afterwards, never drained from a
	// channel: cross-cutting rule 12 means the output must not depend on which goroutine finished first.
	tables := paint{rates: rates, ks: ks, plainM: plainM, plainFloor: plainFloor,
		massifFloor: massifFloor, massifFraction: massifFraction, massifCells: m.Planet.MassifCells(),
		rockCells: m.Planet.LithologyCells(), rockMult: m.Pipeline.Lithology.KMultipliers,
		lithMix: in.Legend.LithologyMixes(), faults: in.Faults,
		// The *manifest's* step count, not the overridden one. `throw_m` is a total displacement over the
		// run and it becomes a rate by dividing by the run's length, so taking the override would make a
		// short run raise the *rate* to build the same scarp in less time - which at `--steps 200` is five
		// times the uplift and past the repose ceiling, so a tuning run would show every fault pinned
		// against the clamp and tell an author nothing about the world they are tuning. `--steps` means
		// "run less time"; everything else in the solve is under-done by it and faults should be too.
		runYears:   float64(m.Pipeline.Fluvial.Steps) * params.DtYr,
		clampCeilM: faultCeilingMYr(m),
	}

	jobs := opt.Jobs
	if jobs <= 0 {
		jobs = 3
	}
	if jobs > len(wanted) {
		jobs = len(wanted)
	}
	// Biggest first, so the long poles start early and the short ones fill the tail.
	order := append([]region.Region(nil), wanted...)
	sort.Slice(order, func(a, b int) bool { return order[a].Cells() > order[b].Cells() })

	out := make([]RegionResult, len(order))
	var wg sync.WaitGroup
	var mu sync.Mutex
	// Guards the composite - the one place workers touch shared state - so that OnRegion can read the whole
	// planet without racing them.
	var composite sync.Mutex
	safeLog := func(format string, a ...any) {
		mu.Lock()
		log(format, a...)
		mu.Unlock()
	}
	next := make(chan int)
	go func() {
		for i := range order {
			next <- i
		}
		close(next)
	}()
	for w := 0; w < jobs; w++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := range next {
				if cancelled(opt.Cancel) {
					continue // drain the queue; the workers already running stop at their next step
				}
				out[i] = solveRegion(in, order[i], params, tables, res, safeLog, opt, &composite)
			}
		}()
	}
	wg.Wait()

	res.Regions = out
	sort.Slice(res.Regions, func(a, b int) bool { return res.Regions[a].ID < res.Regions[b].ID })

	if cancelled(opt.Cancel) {
		res.cancelled = true
		res.Elapsed = time.Since(started)
		log("cancelled after %s", res.Elapsed.Round(time.Second))
		return res, nil
	}

	res.Craters = stampCraters(in, res, log)
	runCoast(in, res, log)
	res.Elapsed = time.Since(started)
	return res, nil
}

func wantedRegions(all []region.Region, only []int) []region.Region {
	if len(only) == 0 {
		return all
	}
	keep := make(map[int]bool, len(only))
	for _, id := range only {
		keep[id] = true
	}
	var out []region.Region
	for _, r := range all {
		if keep[r.ID] {
			out = append(out, r)
		}
	}
	return out
}

// solveRegion is one landmass: cut it out, build its painted geology, run the solve, write the land back.
// progressCells is where a region becomes big enough to be worth reporting on: below a couple of million it
// is over in a minute or two and the lines are noise.
const progressCells = 2_000_000

// paint is the per-class tables a region's geology is built from, bundled so the worker signature does not
// grow a column every time the legend learns a new word.
type paint struct {
	rates, ks   []float32
	plainM      []float64
	plainFloor  []float32
	massifFloor []float32
	// massifFraction is per class, zero for a class that is one rate all over. massifCells is the planet's
	// upland fabric wavelength in lattice cells, which is a property of the planet rather than of a class.
	massifFraction []float64
	massifCells    int

	// The rock field and the fault set: both are properties of the planet that every region reads the same
	// way, which is the whole point of computing them once above rather than per region.
	rockCells  int
	rockMult   []float64
	lithMix    []float64
	faults     []uplift.FaultTrace
	runYears   float64
	clampCeilM float64
}

func cancelled(ch <-chan struct{}) bool {
	if ch == nil {
		return false
	}
	select {
	case <-ch:
		return true
	default:
		return false
	}
}

func solveRegion(in *Inputs, rg region.Region, params fluvial.Params, t paint,
	res *Result, log func(string, ...any), opt BakeOptions, composite *sync.Mutex) RegionResult {

	m := in.M
	start := time.Now()
	class, land := in.Part.Cut(in.Map, rg)

	up := uplift.FromTemplate(uplift.Paint{
		Frame: rg.Frame, Class: class, Land: land,
		Rates: t.rates, Ks: t.ks, PlainM: t.plainM, PlainFloor: t.plainFloor,
		MassifFloor: t.massifFloor, MassifFraction: t.massifFraction, MassifCells: t.massifCells,
		RockCells: t.rockCells, RockMult: t.rockMult, LithMix: t.lithMix,
		Faults: t.faults, RunYears: t.runYears, ClampCeilM: t.clampCeilM,
		Variation: m.Planet.UpliftVariation,
	}, m)

	h := up.Height.Clone()
	grid := fluvial.NewGrid(rg.Frame.W, rg.Frame.H, rg.Frame.P.CellM, up.Base)
	grid.SetSeed(m.Source.Seed)
	grid.SetCancel(opt.Cancel)
	grid.SetFrame(rg.Frame) // the jitter is a hash of world position, not of grid index
	grid.SetElevationRange(m.ElevationM.Min-200, m.ElevationM.Max+500)

	log("region %2d  start   %.1f x %.1f km, %d cells (%d land)%s",
		rg.ID, float64(rg.Frame.W)*rg.Frame.P.CellM/1000, float64(rg.Frame.H)*rg.Frame.P.CellM/1000,
		rg.Cells(), rg.LandCells, seamNote(rg.Seam))

	// Progress, because a big region is an hour on its own and silence for an hour is indistinguishable from
	// a hang. Every tenth, and only for regions big enough to be worth waiting on: with several in flight the
	// lines interleave, so each one carries its region id and none of them is per-step.
	var progress func(step, total int, pct float64)
	if rg.Cells() > progressCells {
		next := 10.0
		progress = func(step, total int, pct float64) {
			if pct < next || step == 0 {
				return
			}
			next = pct + 10
			lo, hi := landExtent(h.Data, land)
			elapsed := time.Since(start)
			eta := time.Duration(float64(elapsed) / (pct / 100) * (1 - pct/100))
			log("region %2d  %3.0f%%    step %d/%d, land %.0f..%.0f m, eta %s",
				rg.ID, pct, step, total, lo, hi, eta.Round(time.Second))
		}
	}
	grid.Run(h.Data, up.Rate.Data, up.K.Data, params, progress)

	// The edge-preserving pass, if the manifest asked for one. Here and not after the write: the region's
	// statistics are gathered a few lines below while the grid is alive, and a smooth applied after them
	// would put a surface on disk that meta.json does not describe - which is the bake-against-tiles drift
	// trap in a new costume.
	//
	// It reads a cell's eight neighbours, and that is safe against the decomposition for a reason particular
	// to this pass: a region's frame edges are open ocean by construction and every landmass in it sits at
	// least margin_cells from them, so a land cell's neighbours are always inside its own region. The rule
	// this would otherwise break - no neighbourhood operation near a region edge - still stands for anything
	// that reaches across the waterline.
	if sm := m.Pipeline.Smooth; sm.Passes > 0 {
		field.SmoothEdgePreserving(h.Data, h.W, h.H, h.CellM, land, sm.Passes, sm.SlopeRef, grid.Scratch())
	}

	lo, hi := landExtent(h.Data, land)
	peakRate := 0.0
	for i, isLand := range land {
		if isLand && float64(up.Rate.Data[i]) > peakRate {
			peakRate = float64(up.Rate.Data[i])
		}
	}
	clip := clipFraction(h.Data, land, m.ElevationM.Min, m.ElevationM.Max)
	secs := time.Since(start).Seconds()
	clampNote := ""
	if up.FaultClamped > 0 {
		clampNote = fmt.Sprintf(", %.2f%% of it bounded at repose by faults",
			100*float64(up.FaultClamped)/float64(rg.LandCells))
	}
	log("region %2d  done    %.0f..%.0f m, %.3f%% clipped, %.2f mm/yr peak%s  [%.0f s]",
		rg.ID, lo, hi, clip*100, peakRate*1000, clampNote, secs)

	// The region's own land statistics, gathered here while its grid is still alive - it is thrown away a
	// few lines below, and the composited planet has no uplift field or flow topology to recover them from.
	// Only the land this region *owns*: `Cut` marks nothing else, so the pieces are disjoint and the sum is
	// the planet.
	acc := stats.New(statsOptions(m))
	acc.Add(stats.Input{
		H: h, Land: land, UpliftMYr: up.Rate.Data, KLocal: up.K.Data,
		Area: grid.Area, Receiver: grid.Receiver, Length: grid.Length,
	})

	rr := RegionResult{ID: rg.ID, Cells: rg.Cells(), LandCells: rg.LandCells,
		Seconds: secs, MinM: lo, MaxM: hi, ClipFrac: clip, FaultClamped: up.FaultClamped, stats: acc}

	// The write-back and the hook under one lock. Composite is the only place a worker touches shared state,
	// so holding it here is what lets OnRegion read the whole planet without racing the others.
	composite.Lock()
	in.Part.Composite(res.Height.Data, in.Map, rg, h.Data)
	compositeFlow(in.Part, in.Map, rg, res.Flow, grid.Area)
	if opt.OnRegion != nil {
		opt.OnRegion(res, rr)
	}
	composite.Unlock()

	return rr
}

func seamNote(seam bool) string {
	if seam {
		return ", across the seam"
	}
	return ""
}

// compositeFlow writes a region's drainage area back for the preview's rivers. It follows Composite's rule:
// only the land this region owns.
func compositeFlow(part *region.Partition, mp *template.Map, rg region.Region, dst, area []float32) {
	for y := 0; y < rg.Frame.H; y++ {
		for x := 0; x < rg.Frame.W; x++ {
			pi := rg.Frame.PlanetIdx(x, y)
			if part.Owner[pi] != int32(rg.ID) || mp.Sea[pi] {
				continue
			}
			dst[pi] = area[y*rg.Frame.W+x]
		}
	}
}

func solveParams(m *manifest.Manifest) fluvial.Params {
	f := m.Pipeline.Fluvial
	return fluvial.Params{
		K: f.K, M: f.M, N: f.N, DtYr: f.DtYr, Steps: f.Steps,
		Diffusion: f.DiffusionM2Yr, FillEvery: f.FillEvery,
		TalusSlope:      thermal.TalusFromDegrees(m.Pipeline.Thermal.TalusDeg),
		ThermalEvery:    m.Pipeline.Thermal.Every,
		ThermalPasses:   m.Pipeline.Thermal.CoarsePasses,
		CriticalAreaM2:  f.CriticalAreaM2,
		ChannelTaper:    f.ChannelTaper,
		CriticalSlope:   thermal.TalusFromDegrees(f.CriticalSlopeDeg),
		SlopeCap:        f.SlopeCap,
		MaxHillslopeSub: f.MaxHillslopeSub,
		MFDExponent:     f.MFDExponent,
	}
}

// runCoast lays the sea floor and works the shoreline, once, over the whole cylinder.
//
// Once and whole rather than per region, which is D-53's rule and is not a convenience: the pass costs tens
// of nanoseconds a cell against tens of nanoseconds a cell *per step* for the solve, and cutting it up would
// truncate the fetch across every strait, split the sediment budget whose conservation is the one thing in it
// not derived from something already measured, and leave the shoreline length and the exposure percentiles as
// statistics that do not pool. It runs after every region is composited, because two of its three processes
// read the finished land: the shelf width comes off the relief standing behind each shore, and the surf cuts
// into whatever the solve built.
//
// The painted depths go in as the *abyss*, one value per cell. That is what makes the derived margin meet the
// painted ocean instead of stepping to it: an author who painted `shelf` at 120 m, `deep` at 512 and `surf`
// at 20 gets a continental slope that runs down to each of those where each of them is, and a strait painted
// shallower than the shelf break comes out as shelf all the way across rather than as a trench.
func runCoast(in *Inputs, res *Result, log func(string, ...any)) {
	m := in.M
	depths := in.Legend.Depths()
	abyss := make([]float32, len(res.Height.Data))
	for i := range abyss {
		if in.Map.Sea[i] {
			abyss[i] = depths[in.Map.Class[i]]
		}
	}

	start := time.Now()
	cs := coast.Build(coast.Input{
		Height: res.Height, Sea: in.Map.Sea, SeaLevelM: m.SeaLevelM,
		BreakM: m.ShelfBreakM(),
		// AbyssM is the square canvas's key and is never read on this path: `Abyss` below is per-cell and
		// always allocated, so `abyssAt` takes the painting every time. Left as the fallback it is rather
		// than removed, and named here because it is the same shape as the read D-64 had to fix one line up.
		AbyssM: -m.Pipeline.Continent.SeaFloorM.Lo(),
		Abyss:  abyss,
		// The cylinder, which is the whole of what D-59's successor had to add: every march, every ray and
		// every running sum in the pass crosses the seam now, and a planet measured on a flat grid would have
		// had its shelf, its fetch and its sediment budget all stop dead at one meridian.
		WrapX:        true,
		NoisePeriodM: in.P.NoisePeriodM,
		Flow:         res.Flow,
		Seed:         m.Source.Seed,
		Cfg:          m.Pipeline.Coast,
	})
	res.Coast = cs
	res.Sea = cs.Sea

	if !m.Pipeline.Coast.Enabled {
		log("ocean   the painted depths, flat: the coastal pass is switched off in the manifest")
		return
	}
	// The margin's own numbers, said out loud, because they decide how much of the painted ocean survives
	// and they were a square-canvas inheritance nobody could see until D-64. A shelf and a slope together
	// reach `ShelfKm.Hi() + SlopeKm` from every shore; where the sea is narrower than twice that, the
	// painting's depth is never reached anywhere in it and the author's ocean is whatever `break_m` says.
	cfg := m.Pipeline.Coast
	reach := cfg.ShelfKm.Hi() + cfg.SlopeKm
	log("        sea floor: shelf %.1f..%.1f km to a break at %.0f m, then %.1f km of slope to the painted "+
		"depth, so the painting owns the water past %.1f km offshore and nothing nearer",
		cfg.ShelfKm.Lo(), cfg.ShelfKm.Hi(), m.ShelfBreakM(), cfg.SlopeKm, reach)
	log("coast   %s", cs.Stats.Summary())
	log("        over the whole cylinder in %s, once: the pass is tens of nanoseconds a cell and cutting it",
		time.Since(start).Round(time.Millisecond))
	log("        up would truncate the fetch across every strait and split the sediment budget")
}

func landExtent(h []float32, land []bool) (lo, hi float64) {
	lo, hi = 1e30, -1e30
	any := false
	for i, v := range h {
		if !land[i] {
			continue
		}
		any = true
		if float64(v) < lo {
			lo = float64(v)
		}
		if float64(v) > hi {
			hi = float64(v)
		}
	}
	if !any {
		return 0, 0
	}
	return lo, hi
}

// clipFraction is how much of the land the 16-bit encoding would cut off. Above a fraction of a per cent it
// is a failed run rather than a rounded one, and painted uplift makes it easier to hit: an author can ask
// for more relief than the elevation range holds.
func clipFraction(h []float32, land []bool, minM, maxM float64) float64 {
	n, clipped := 0, 0
	for i, v := range h {
		if !land[i] {
			continue
		}
		n++
		if float64(v) < minM || float64(v) > maxM {
			clipped++
		}
	}
	if n == 0 {
		return 0
	}
	return float64(clipped) / float64(n)
}

// faultCeilingMYr is the uplift rate at which a divide stands at the angle of repose, in metres a year at
// K x1. It is the number `terrain plan` already prints as the clamp ceiling, in the unit the solve works in;
// here it bounds what a *fault* may add on top of what an author painted, and nothing else. Zero when the
// repose clamp is switched off, which switches the bound off with it.
func faultCeilingMYr(m *manifest.Manifest) float64 {
	if m.Pipeline.Thermal.TalusDeg <= 0 || m.Pipeline.Thermal.TalusDeg >= 90 {
		return 0
	}
	return clampCeiling(m.Pipeline.Fluvial.K, m.GeologyCellM(), m.Pipeline.Fluvial.M,
		m.Pipeline.Thermal.TalusDeg) / 1000
}