package main

import (
	"fmt"
	"sort"
	"strconv"
	"strings"
	"unicode"
)

type ShapeDefinition struct {
	Name  string `json:"name"`
	Frets []int  `json:"frets"` // -1 = muted
}

type ShapeQuery struct {
	Shape         ShapeDefinition `json:"shape"`
	TargetQuality string          `json:"target_quality"`
	TargetRoot    int             `json:"target_root"` // 0-11 or -1 for any
	Voicing       []string        `json:"voicing"`     // per-string "" or "C4"/"60"
}

type CompanionChord struct {
	Shape   string   `json:"shape"`
	Chord   string   `json:"chord"`
	Root    string   `json:"root"`
	Quality string   `json:"quality"`
	Notes   []string `json:"notes"`
}

type TuningCandidate struct {
	Tuning      []string         `json:"tuning"`
	TuningMIDI  []int            `json:"tuning_midi"`
	Root        string           `json:"root"`
	Chord       string           `json:"chord"`
	Companions  []CompanionChord `json:"companions"`
	ValidChords int              `json:"valid_chords"`
	MajMinCount int              `json:"maj_min_count"`
	HighCompat  bool             `json:"high_compat"`
}

var standardMIDI = []int{40, 45, 50, 55, 59, 64} // E2 A2 D3 G3 B3 E4

func DefaultShapes() []ShapeDefinition {
	return []ShapeDefinition{
		{"E major", []int{0, 2, 2, 1, 0, 0}},
		{"E minor", []int{0, 2, 2, 0, 0, 0}},
		{"A major", []int{-1, 0, 2, 2, 2, 0}},
		{"A minor", []int{-1, 0, 2, 2, 1, 0}},
		{"D major", []int{-1, -1, 0, 2, 3, 2}},
		{"D minor", []int{-1, -1, 0, 2, 3, 1}},
		{"E7 (Dom7)", []int{0, 2, 0, 1, 0, 0}},
		{"Em7 (Min7)", []int{0, 2, 0, 0, 0, 0}},
		{"A7 (Dom7)", []int{-1, 0, 2, 0, 2, 0}},
		{"Am7 (Min7)", []int{-1, 0, 2, 0, 1, 0}},
		{"Lazy barre", []int{0, 0, 2, 2, 2, 2}},
	}
}

func parsePitchInput(s string) (midi int, pc int, err error) {
	s = strings.TrimSpace(s)
	if s == "" {
		return -1, -1, fmt.Errorf("empty input")
	}

	if n, e := strconv.Atoi(s); e == nil {
		if n < 0 || n > 127 {
			return 0, 0, fmt.Errorf("MIDI %d out of range", n)
		}
		return n, n % 12, nil
	}

	// Note name: C4, Bb3, F#5, etc.
	i := 0
	if i >= len(s) {
		return 0, 0, fmt.Errorf("invalid note: %s", s)
	}
	notePart := string(unicode.ToUpper(rune(s[i])))
	i++
	for i < len(s) && (s[i] == '#' || s[i] == 'b') {
		notePart += string(s[i])
		i++
	}

	sem, ok := NoteToSemitone[notePart]
	if !ok {
		return 0, 0, fmt.Errorf("unknown note: %s", notePart)
	}

	if i >= len(s) {
		return 0, 0, fmt.Errorf("missing octave in: %s", s)
	}

	octStr := s[i:]
	neg := false
	if octStr[0] == '-' {
		neg = true
		octStr = octStr[1:]
	}
	oct, e := strconv.Atoi(octStr)
	if e != nil {
		return 0, 0, fmt.Errorf("invalid octave in: %s", s)
	}
	if neg {
		oct = -oct
	}

	midi = (oct+1)*12 + sem
	if midi < 0 || midi > 127 {
		return 0, 0, fmt.Errorf("MIDI %d out of range for %s", midi, s)
	}
	return midi, sem, nil
}

func closestMIDIInRange(pc, standard int) (int, bool) {
	lo := standard - 7
	hi := standard + 7
	// find MIDI note with pitch class pc closest to standard within range
	best := -1
	bestDist := 999
	for m := lo; m <= hi; m++ {
		if m < 0 || m > 127 {
			continue
		}
		if m%12 == pc {
			d := m - standard
			if d < 0 {
				d = -d
			}
			if d < bestDist {
				bestDist = d
				best = m
			}
		}
	}
	return best, best >= 0
}

func findTuningsForShape(query ShapeQuery, allShapes []ShapeDefinition) ([]TuningCandidate, error) {
	shape := query.Shape.Frets
	nStrings := len(shape)
	if nStrings != 6 {
		return nil, fmt.Errorf("shape must have 6 strings")
	}

	// resolve target intervals
	var targetIntervals []int
	defs := GetChordDefinitions()
	found := false
	for _, cat := range defs {
		if ivs, ok := cat[query.TargetQuality]; ok {
			targetIntervals = ivs
			found = true
			break
		}
	}
	if !found {
		return nil, fmt.Errorf("unknown quality: %s", query.TargetQuality)
	}

	// parse voicing pins
	type pinInfo struct {
		midi int
		pc   int
	}
	pins := make([]pinInfo, nStrings)
	for i := range pins {
		pins[i] = pinInfo{-1, -1}
	}
	if len(query.Voicing) > 0 {
		for i := 0; i < nStrings && i < len(query.Voicing); i++ {
			v := strings.TrimSpace(query.Voicing[i])
			if v == "" {
				continue
			}
			if shape[i] == -1 {
				continue // muted string, ignore voicing
			}
			m, p, err := parsePitchInput(v)
			if err != nil {
				return nil, fmt.Errorf("string %d voicing: %v", i+1, err)
			}
			pins[i] = pinInfo{m, p}
		}
	}

	// determine candidate roots
	var roots []int
	if query.TargetRoot >= 0 && query.TargetRoot < 12 {
		roots = []int{query.TargetRoot}
	} else {
		roots = make([]int, 12)
		for i := 0; i < 12; i++ {
			roots[i] = i
		}
	}

	// chord pitch classes for each root
	type tuningResult struct {
		tuningPCs  []int
		tuningMIDI []int
		root       int
	}
	var results []tuningResult
	seen := make(map[string]bool)

	for _, root := range roots {
		chordPCs := make(map[int]bool, len(targetIntervals))
		for _, iv := range targetIntervals {
			chordPCs[(root+iv)%12] = true
		}

		// for each non-muted string, compute candidate open note PCs
		type stringCandidates struct {
			pcs []int
		}
		candidates := make([]stringCandidates, nStrings)
		muted := make([]bool, nStrings)

		for s := 0; s < nStrings; s++ {
			if shape[s] == -1 {
				muted[s] = true
				candidates[s] = stringCandidates{[]int{standardMIDI[s] % 12}}
				continue
			}

			if pins[s].pc >= 0 {
				// pinned: open note = pinned_midi - shape_fret
				openPC := (pins[s].pc - shape[s] + 120) % 12
				candidates[s] = stringCandidates{[]int{openPC}}
				continue
			}

			var cands []int
			for ct := range chordPCs {
				openPC := (ct - shape[s] + 12) % 12
				cands = append(cands, openPC)
			}
			candidates[s] = stringCandidates{cands}
		}

		// cartesian product
		indices := make([]int, nStrings)
		sizes := make([]int, nStrings)
		total := 1
		for s := 0; s < nStrings; s++ {
			sizes[s] = len(candidates[s].pcs)
			total *= sizes[s]
		}

		for combo := 0; combo < total; combo++ {
			tmp := combo
			for s := nStrings - 1; s >= 0; s-- {
				indices[s] = tmp % sizes[s]
				tmp /= sizes[s]
			}

			tuningPCs := make([]int, nStrings)
			for s := 0; s < nStrings; s++ {
				tuningPCs[s] = candidates[s].pcs[indices[s]]
			}

			// check all chord tones present in voicing (non-muted strings)
			voicedPCs := make(map[int]bool)
			for s := 0; s < nStrings; s++ {
				if !muted[s] {
					sounded := (tuningPCs[s] + shape[s]) % 12
					voicedPCs[sounded] = true
				}
			}
			allPresent := true
			for ct := range chordPCs {
				if !voicedPCs[ct] {
					allPresent = false
					break
				}
			}
			if !allPresent {
				continue
			}

			// resolve MIDI pitches with tension check
			tuningMIDI := make([]int, nStrings)
			valid := true
			for s := 0; s < nStrings; s++ {
				if pins[s].midi >= 0 && !muted[s] {
					openMIDI := pins[s].midi - shape[s]
					if openMIDI < 0 || openMIDI > 127 {
						valid = false
						break
					}
					if openMIDI%12 != tuningPCs[s] {
						valid = false
						break
					}
					std := standardMIDI[s]
					diff := openMIDI - std
					if diff < 0 {
						diff = -diff
					}
					if diff > 7 {
						valid = false
						break
					}
					tuningMIDI[s] = openMIDI
				} else {
					m, ok := closestMIDIInRange(tuningPCs[s], standardMIDI[s])
					if !ok {
						valid = false
						break
					}
					tuningMIDI[s] = m
				}
			}
			if !valid {
				continue
			}

			key := fmt.Sprint(tuningMIDI)
			if !seen[key] {
				seen[key] = true
				results = append(results, tuningResult{tuningPCs, tuningMIDI, root})
			}
		}
	}

	deduped := results

	// build output
	var output []TuningCandidate
	for _, r := range deduped {
		tuningNames := make([]string, nStrings)
		for s := 0; s < nStrings; s++ {
			tuningNames[s] = midiToNoteName(r.tuningMIDI[s])
		}

		companions := identifyCompanions(r.tuningMIDI, allShapes, query.Shape.Name, r.root, query.TargetQuality)

		validChords := 0
		majMinCount := 0
		for _, c := range companions {
			if c.Quality != "" && c.Chord != "?" && c.Chord != "—" {
				validChords++
			}
			if c.Quality == "major" || c.Quality == "minor" {
				majMinCount++
			}
		}
		totalShapes := len(companions)
		threshold := totalShapes / 2
		if threshold < 5 {
			threshold = 5
		}

		output = append(output, TuningCandidate{
			Tuning:      tuningNames,
			TuningMIDI:  r.tuningMIDI,
			Root:        SemitoneToNote[r.root],
			Chord:       SemitoneToNote[r.root] + " " + query.TargetQuality,
			Companions:  companions,
			ValidChords: validChords,
			MajMinCount: majMinCount,
			HighCompat:  majMinCount > threshold,
		})
	}

	sort.Slice(output, func(i, j int) bool {
		if output[i].ValidChords != output[j].ValidChords {
			return output[i].ValidChords > output[j].ValidChords
		}
		if output[i].MajMinCount != output[j].MajMinCount {
			return output[i].MajMinCount > output[j].MajMinCount
		}
		return false
	})

	return output, nil
}

func identifyCompanions(tuningMIDI []int, shapes []ShapeDefinition, searchShapeName string, targetRoot int, targetQuality string) []CompanionChord {
	defs := GetChordDefinitions()
	nStrings := len(tuningMIDI)

	type chordDef struct {
		quality   string
		intervals []int
	}
	var allDefs []chordDef
	for _, cat := range defs {
		for name, ivs := range cat {
			allDefs = append(allDefs, chordDef{name, ivs})
		}
	}

	// put search shape first
	ordered := make([]ShapeDefinition, 0, len(shapes))
	searchIdx := -1
	for i, s := range shapes {
		if s.Name == searchShapeName {
			searchIdx = i
			break
		}
	}
	if searchIdx >= 0 {
		ordered = append(ordered, shapes[searchIdx])
		for i, s := range shapes {
			if i != searchIdx {
				ordered = append(ordered, s)
			}
		}
	} else {
		ordered = shapes
	}

	var companions []CompanionChord
	for _, shape := range ordered {
		frets := shape.Frets
		if len(frets) != nStrings {
			continue
		}

		var soundedPCs []int
		notes := make([]string, nStrings)
		allMuted := true
		for s := 0; s < nStrings; s++ {
			if frets[s] == -1 {
				notes[s] = "x"
				continue
			}
			allMuted = false
			midi := tuningMIDI[s] + frets[s]
			pc := midi % 12
			soundedPCs = append(soundedPCs, pc)
			notes[s] = midiToNoteName(midi)
		}

		if allMuted || len(soundedPCs) == 0 {
			companions = append(companions, CompanionChord{
				Shape: shape.Name,
				Chord: "—",
				Notes: notes,
			})
			continue
		}

		// search shape: we already know the target chord
		if shape.Name == searchShapeName && targetRoot >= 0 {
			rn := SemitoneToNote[targetRoot]
			companions = append(companions, CompanionChord{
				Shape:   shape.Name,
				Chord:   rn + " " + targetQuality,
				Root:    rn,
				Quality: targetQuality,
				Notes:   notes,
			})
			continue
		}

		uniquePCs := uniqueInts(soundedPCs)

		// collect all valid interpretations, pick the simplest
		type candidate struct {
			root    int
			quality string
			size    int
			bassIdx int
		}
		var candidates []candidate
		for ri, root := range uniquePCs {
			ivSet := make(map[int]bool)
			for _, pc := range soundedPCs {
				ivSet[(pc-root+12)%12] = true
			}
			for _, cd := range allDefs {
				if len(ivSet) != len(cd.intervals) {
					continue
				}
				ok := true
				for _, iv := range cd.intervals {
					if !ivSet[iv] {
						ok = false
						break
					}
				}
				if ok {
					candidates = append(candidates, candidate{root, cd.quality, len(cd.intervals), ri})
				}
			}
		}

		bestChord := "?"
		bestRoot := ""
		bestQuality := ""
		if len(candidates) > 0 {
			best := candidates[0]
			for _, c := range candidates[1:] {
				if c.size < best.size || (c.size == best.size && c.bassIdx < best.bassIdx) {
					best = c
				}
			}
			rn := SemitoneToNote[best.root]
			bestChord = rn + " " + best.quality
			bestRoot = rn
			bestQuality = best.quality
		}

		companions = append(companions, CompanionChord{
			Shape:   shape.Name,
			Chord:   bestChord,
			Root:    bestRoot,
			Quality: bestQuality,
			Notes:   notes,
		})
	}

	return companions
}

func midiToNoteName(midi int) string {
	pc := midi % 12
	octave := midi/12 - 1
	return fmt.Sprintf("%s%d", SemitoneToNote[pc], octave)
}