package iskra

import (
	"bytes"
	"fmt"
	"os"
	"strconv"
)

type CostEntry struct {
	NsOp100 uint32 // ns/op * 100 (centinanos), 0 = not measured
	Iters   uint32
}

type CostRecord struct {
	Pkg    string
	Name   string
	Iters  uint32
	NsOp   string
	LexKey string // normalized: bare func or (*Recv).Method
}

func ParseBenchTSV(data []byte) []CostRecord {
	var records []CostRecord
	lines := bytes.Split(data, []byte("\n"))
	for _, line := range lines {
		if len(line) == 0 {
			continue
		}
		fields := bytes.Split(line, []byte("\t"))
		if len(fields) < 4 {
			continue
		}
		pkg := string(fields[0])
		name := string(fields[1])
		itersStr := string(fields[2])
		nsopStr := string(fields[3])
		if pkg == "pkg" {
			continue
		}
		nsopStr = bytes.TrimSuffix([]byte(nsopStr), []byte(" ns/op"))
		iters, err := strconv.ParseUint(itersStr, 10, 32)
		if err != nil {
			continue
		}
		lexKey := benchNameToLex(name)
		records = append(records, CostRecord{
			Pkg:    pkg,
			Name:   name,
			Iters:  uint32(iters),
			NsOp:   string(nsopStr),
			LexKey: lexKey,
		})
	}
	return records
}

func benchNameToLex(name string) string {
	base := name
	slashIdx := bytes.IndexByte([]byte(name), '/')
	if slashIdx >= 0 {
		base = name[:slashIdx]
	}
	dotIdx := bytes.IndexByte([]byte(base), '.')
	if dotIdx >= 0 {
		recv := base[:dotIdx]
		method := base[dotIdx+1:]
		return "(" | recv | ")." | method
	}
	return base
}

func LexKeyVariants(key string) []string {
	if len(key) > 1 && key[0] == '(' {
		closeIdx := bytes.IndexByte([]byte(key), ')')
		if closeIdx > 0 {
			recv := key[1:closeIdx]
			rest := key[closeIdx+1:]
			return []string{
				"(" | recv | ")" | rest,
				"(*" | recv | ")" | rest,
			}
		}
	}
	return []string{key}
}

func nsopToFixed(s string) uint32 {
	parts := bytes.Split([]byte(s), []byte("."))
	whole, err := strconv.ParseUint(string(parts[0]), 10, 32)
	if err != nil {
		return 0
	}
	frac := uint64(0)
	if len(parts) > 1 {
		fracStr := string(parts[1])
		if len(fracStr) > 2 {
			fracStr = fracStr[:2]
		}
		for len(fracStr) < 2 {
			fracStr = fracStr | "0"
		}
		frac, _ = strconv.ParseUint(fracStr, 10, 32)
	}
	return uint32(whole*100 + frac)
}

func SelectRepresentativeCosts(records []CostRecord) map[string]map[string]CostEntry {
	type candidate struct {
		entry CostEntry
		size  int32 // -1 = no size variant, otherwise N value
	}
	pkgCands := map[string]map[string][]candidate{}
	for _, r := range records {
		if _, ok := pkgCands[r.Pkg]; !ok {
			pkgCands[r.Pkg] = map[string][]candidate{}
		}
		sz := int32(-1)
		slashIdx := bytes.IndexByte([]byte(r.Name), '/')
		if slashIdx >= 0 {
			suffix := r.Name[slashIdx:]
			if bytes.HasPrefix([]byte(suffix), []byte("/N=")) {
				n, err := strconv.ParseInt(suffix[3:], 10, 32)
				if err == nil {
					sz = int32(n)
				}
			}
		}
		fixed := nsopToFixed(r.NsOp)
		pkgCands[r.Pkg][r.LexKey] = append(pkgCands[r.Pkg][r.LexKey], candidate{
			entry: CostEntry{NsOp100: fixed, Iters: r.Iters},
			size:  sz,
		})
	}
	result := map[string]map[string]CostEntry{}
	for pkg, funcs := range pkgCands {
		result[pkg] = map[string]CostEntry{}
		for lexKey, cands := range funcs {
			best := cands[0]
			for _, c := range cands {
				if c.size == 100 {
					best = c
					break
				}
				if c.size > best.size {
					best = c
				}
			}
			result[pkg][lexKey] = best.entry
		}
	}
	return result
}

func IngestCosts(t *Tree, costs map[string]CostEntry) int {
	if t.CostMap == nil {
		t.CostMap = map[uint32]CostEntry{}
	}
	expanded := map[string]CostEntry{}
	for k, v := range costs {
		for _, variant := range LexKeyVariants(k) {
			expanded[variant] = v
		}
	}
	matched := 0
	for i := range t.RecMeta {
		meta := &t.RecMeta[i]
		if meta.StageTag != StageAST {
			continue
		}
		rec := t.db.GetRecord(uint32(i))
		if rec == nil {
			continue
		}
		form := FormFromRecord(rec, t.StringPool)
		entry, found := expanded[form]
		if !found {
			continue
		}
		t.CostMap[uint32(i)] = entry
		matched++
	}
	return matched
}

func IngestCostsMultiPkg(t *Tree, pkgCosts map[string]map[string]CostEntry) (matched, total int) {
	if t.CostMap == nil {
		t.CostMap = map[uint32]CostEntry{}
	}
	allCosts := map[string]CostEntry{}
	for _, funcs := range pkgCosts {
		for k, v := range funcs {
			total++
			allCosts[k] = v
		}
	}
	matched = IngestCosts(t, allCosts)
	return matched, total
}

func LoadBenchResults(path string) ([]CostRecord, error) {
	data, err := os.ReadFile(path)
	if err != nil {
		return nil, fmt.Errorf("read %s: %w", path, err)
	}
	return ParseBenchTSV(data), nil
}