package main

// theme.go turns a set of scanned Tailwind tokens into resolved colours by driving
// kjol's own Tailwind engine (package tw) and reading back what it emits. We do NOT
// re-implement utility parsing: we hand the engine every candidate, let it compile,
// and then read the CSS it produced. That keeps aria-check faithful to whatever the
// real build does — opacity modifiers, arbitrary values, semantic tokens, the lot —
// and correct-by-construction as the engine evolves.
//
// The engine gives us three things in its output:
//
//   - the `:root` custom-property block  → the LIGHT variable environment
//   - the `.dark { … }` override block   → the DARK variable environment (overlay)
//   - the `@layer utilities` rules        → token → colour-valued declaration
//
// A token's colour is then just: look up its declaration's value expression, and
// resolve it (var() chains and color-mix()) against the chosen environment.

import (
	"regexp"
	"strconv"
	"strings"

	"kjol/tw"
)

// Resolver holds everything needed to turn a token into a concrete colour in either
// theme.
type Resolver struct {
	light map[string]string // --var → value expression, light theme
	dark  map[string]string // --var → value expression, dark theme (light overlaid)

	// token → the colour declaration the engine emitted for it. prop is "color"
	// (a text-* utility) or "background-color" (a bg-* utility); expr is the raw
	// value, e.g. "var(--color-ink)" or "color-mix(in oklab, var(--color-red-500) 50%, transparent)".
	tokens map[string]tokenDecl
}

type tokenDecl struct {
	prop string
	expr string
}

// surfaceExpr is the page background token; a translucent background composites onto
// it (see check.go). It is a plain --var lookup in whichever environment.
const surfaceVar = "--color-surface"

// NewResolver compiles candidates through the kjol Tailwind engine and indexes the
// result. entryCSS is the app's brand stylesheet (may be empty — kjol's own theme
// layer is always included via tw.CompileApp); baseDir is what any @import/@source in
// the entry resolves against.
func NewResolver(entryCSS, baseDir string, candidates []string) (*Resolver, error) {
	if strings.TrimSpace(entryCSS) == "" {
		entryCSS = "@theme {}"
	}
	css, _, err := tw.CompileApp(entryCSS, baseDir, candidates)
	if err != nil {
		return nil, err
	}
	r := &Resolver{
		light:  map[string]string{},
		dark:   map[string]string{},
		tokens: map[string]tokenDecl{},
	}
	r.indexVars(css)
	r.indexUtilities(css)
	return r, nil
}

// reVarDecl matches a `--name: value;` custom-property declaration.
var reVarDecl = regexp.MustCompile(`(--[A-Za-z0-9-]+)\s*:\s*([^;]+);`)

// indexVars reads the theme `:root`/`:host` block into the light environment and the
// top-level `.dark { … }` rule into the dark overlay (which starts as a copy of
// light). The `.dark` rule we want is the design system's token override — selector
// exactly `.dark`, not the escaped utility selectors like `.dark\:bg-surface`.
func (r *Resolver) indexVars(css string) {
	// Light: every custom property declared under the theme layer's :root/:host.
	// The engine emits all theme variables there (it does not prune unused ones),
	// so a single pass over the block captures the whole palette + tokens.
	if root := blockBody(css, `:root, :host {`); root != "" {
		for _, m := range reVarDecl.FindAllStringSubmatch(root, -1) {
			r.light[m[1]] = strings.TrimSpace(m[2])
		}
	}
	// Some variables (e.g. the FA style flags) sit in a plain `:root {` the engine
	// passes through; fold those in too so nothing referenced dangles.
	if root := blockBody(css, "\n:root {"); root != "" {
		for _, m := range reVarDecl.FindAllStringSubmatch(root, -1) {
			if _, ok := r.light[m[1]]; !ok {
				r.light[m[1]] = strings.TrimSpace(m[2])
			}
		}
	}

	// Dark starts as a copy of light, then every top-level `.dark { … }` rule
	// re-points a subset — the kjol design-system layer defines one, and an app's
	// brand stylesheet may add more, so all of them are folded in, in order.
	for k, v := range r.light {
		r.dark[k] = v
	}
	for _, darkBody := range eachBlock(css, "\n.dark {") {
		for _, m := range reVarDecl.FindAllStringSubmatch(darkBody, -1) {
			r.dark[m[1]] = strings.TrimSpace(m[2])
		}
	}
}

// reColorDecl finds the first color / background-color declaration in a rule body,
// even when it is nested inside a variant wrapper (`&:where(.dark, …) { … }`).
var reColorDecl = regexp.MustCompile(`(?:^|[{\s])(background-color|color)\s*:\s*([^;]+);`)

// reUtilitySelector matches the start of one top-level utility rule and captures its
// (still CSS-escaped) selector, e.g. `.dark\:bg-surface {`.
var reUtilitySelector = regexp.MustCompile(`(?m)^\s{2}\.([^\s{]+)\s*\{`)

// indexUtilities walks the @layer utilities block and records, per token, the first
// colour declaration the engine produced for it. Tokens with no colour declaration
// (layout utilities, font sizes, …) are simply absent from the map — which is
// exactly how we tell a colour utility from a non-colour one.
func (r *Resolver) indexUtilities(css string) {
	body := blockBody(css, "@layer utilities {")
	if body == "" {
		return
	}
	locs := reUtilitySelector.FindAllStringSubmatchIndex(body, -1)
	for i, loc := range locs {
		escSel := body[loc[2]:loc[3]]
		// The rule body runs from this selector's opening brace to the next
		// top-level rule (or the end of the layer). That span may contain nested
		// braces; we only need the first colour declaration within it.
		start := loc[1]
		end := len(body)
		if i+1 < len(locs) {
			end = locs[i+1][0]
		}
		rule := body[start:end]
		m := reColorDecl.FindStringSubmatch(rule)
		if m == nil {
			continue
		}
		token := unescapeIdent(escSel)
		r.tokens[token] = tokenDecl{prop: m[1], expr: strings.TrimSpace(m[2])}
	}
}

// blockBody returns the text between the braces of the first block whose header
// (including its opening `{`) matches marker. It is brace-aware, so nested rules are
// returned intact.
func blockBody(css, marker string) string {
	idx := strings.Index(css, marker)
	if idx < 0 {
		return ""
	}
	open := idx + len(marker) - 1 // position of the '{' in the marker
	depth := 0
	for i := open; i < len(css); i++ {
		switch css[i] {
		case '{':
			depth++
		case '}':
			depth--
			if depth == 0 {
				return css[open+1 : i]
			}
		}
	}
	return ""
}

// eachBlock returns the bodies of every block whose header matches marker, in order.
func eachBlock(css, marker string) []string {
	var out []string
	for {
		idx := strings.Index(css, marker)
		if idx < 0 {
			return out
		}
		body := blockBody(css[idx:], marker)
		out = append(out, body)
		// Advance past this block's opening brace to find the next match.
		css = css[idx+len(marker):]
	}
}

// unescapeIdent reverses CSS identifier escaping so a compiled selector maps back to
// the token the scanner saw. It handles both backslash-escaped punctuation
// (`bg-\[\#fff\]` → `bg-[#fff]`) and numeric escapes (`\32 xl` → `2xl`).
func unescapeIdent(s string) string {
	var b strings.Builder
	for i := 0; i < len(s); i++ {
		if s[i] != '\\' || i+1 >= len(s) {
			b.WriteByte(s[i])
			continue
		}
		i++
		// Numeric escape: 1–6 hex digits, optional single trailing space.
		if isHex(s[i]) {
			j := i
			for j < len(s) && j-i < 6 && isHex(s[j]) {
				j++
			}
			var code int
			for k := i; k < j; k++ {
				code = code*16 + hexVal(s[k])
			}
			if j < len(s) && s[j] == ' ' {
				j++
			}
			b.WriteRune(rune(code))
			i = j - 1
			continue
		}
		b.WriteByte(s[i])
	}
	return b.String()
}

func isHex(c byte) bool {
	return (c >= '0' && c <= '9') || (c >= 'a' && c <= 'f') || (c >= 'A' && c <= 'F')
}

func hexVal(c byte) int {
	switch {
	case c >= '0' && c <= '9':
		return int(c - '0')
	case c >= 'a' && c <= 'f':
		return int(c-'a') + 10
	default:
		return int(c-'A') + 10
	}
}

// Colour resolution ---------------------------------------------------------

// theme selects which variable environment a resolution runs against.
type theme int

const (
	light theme = iota
	dark
)

func (r *Resolver) env(t theme) map[string]string {
	if t == dark {
		return r.dark
	}
	return r.light
}

// isColorToken reports whether a token compiled to a colour-valued text-*/bg-*
// utility, and which side it lands on. side is "fg" for a text colour, "bg" for a
// background colour, "" if the token is not a foreground/background colour utility.
func (r *Resolver) side(token string) string {
	d, ok := r.tokens[token]
	if !ok {
		return ""
	}
	switch d.prop {
	case "color":
		return "fg"
	case "background-color":
		return "bg"
	}
	return ""
}

// resolveToken resolves a scanned token to a colour in the given theme. ok=false
// means the token is not a resolvable colour (unknown, or currentcolor/inherit).
func (r *Resolver) resolveToken(token string, t theme) (RGBA, bool) {
	d, ok := r.tokens[token]
	if !ok {
		return RGBA{}, false
	}
	return r.resolveExpr(d.expr, t, 0)
}

// surface returns the page background colour for a theme — the backdrop a
// translucent background is flattened against.
func (r *Resolver) surface(t theme) (RGBA, bool) {
	if v, ok := r.env(t)[surfaceVar]; ok {
		return r.resolveExpr(v, t, 0)
	}
	return RGBA{}, false
}

var reVarFn = regexp.MustCompile(`^var\(\s*(--[A-Za-z0-9-]+)\s*(?:,\s*([^)]*))?\)$`)

// resolveExpr resolves a CSS colour value expression to an RGBA. It follows var()
// chains through the environment and evaluates the color-mix() form the engine emits
// for opacity; anything else is handed to the literal parser. depth guards against a
// pathological variable cycle.
func (r *Resolver) resolveExpr(expr string, t theme, depth int) (RGBA, bool) {
	expr = strings.TrimSpace(expr)
	if depth > 32 || expr == "" {
		return RGBA{}, false
	}
	if strings.HasPrefix(expr, "var(") {
		m := reVarFn.FindStringSubmatch(expr)
		if m == nil {
			return RGBA{}, false
		}
		if v, ok := r.env(t)[m[1]]; ok {
			return r.resolveExpr(v, t, depth+1)
		}
		if m[2] != "" { // var() fallback
			return r.resolveExpr(m[2], t, depth+1)
		}
		return RGBA{}, false
	}
	if strings.HasPrefix(expr, "color-mix(") {
		return r.resolveColorMix(expr, t, depth)
	}
	return parseLiteralColor(expr)
}

// resolveColorMix evaluates `color-mix(in <space>, <c1> [p1%], <c2> [p2%])`. The
// mixing space in Tailwind's output is always oklab; we evaluate there. This covers
// the opacity form (`… <color> P%, transparent`) and hand-written arbitrary mixes.
func (r *Resolver) resolveColorMix(expr string, t theme, depth int) (RGBA, bool) {
	inner := expr[strings.IndexByte(expr, '(')+1 : strings.LastIndexByte(expr, ')')]
	parts := splitTopLevel(inner, ',')
	if len(parts) != 3 {
		return RGBA{}, false
	}
	// parts[0] is "in oklab" (or another space) — we always mix in oklab.
	c1, w1, ok1 := r.mixComponent(parts[1], t, depth)
	c2, w2, ok2 := r.mixComponent(parts[2], t, depth)
	if !ok1 || !ok2 {
		return RGBA{}, false
	}
	// If only one side gave a percentage, the other takes the remainder.
	if w1 < 0 && w2 < 0 {
		w1, w2 = 0.5, 0.5
	} else if w1 < 0 {
		w1 = clamp01(1 - w2)
	} else if w2 < 0 {
		w2 = clamp01(1 - w1)
	}
	return mixOKLab(c1, w1, c2, w2), true
}

// mixComponent parses one "<color> [P%]" argument of a color-mix(). A negative
// weight means no percentage was given.
func (r *Resolver) mixComponent(s string, t theme, depth int) (RGBA, float64, bool) {
	s = strings.TrimSpace(s)
	weight := -1.0
	if i := strings.LastIndexByte(s, ' '); i >= 0 && strings.HasSuffix(s, "%") {
		if v, err := strconv.ParseFloat(strings.TrimSuffix(s[i+1:], "%"), 64); err == nil {
			weight = v / 100
			s = strings.TrimSpace(s[:i])
		}
	}
	c, ok := r.resolveExpr(s, t, depth+1)
	return c, weight, ok
}

// splitTopLevel splits s on sep, ignoring separators nested inside parentheses.
func splitTopLevel(s string, sep byte) []string {
	var out []string
	depth, start := 0, 0
	for i := 0; i < len(s); i++ {
		switch s[i] {
		case '(':
			depth++
		case ')':
			depth--
		case sep:
			if depth == 0 {
				out = append(out, strings.TrimSpace(s[start:i]))
				start = i + 1
			}
		}
	}
	out = append(out, strings.TrimSpace(s[start:]))
	return out
}