package main

import (
	"git.smesh.lol/moxie/pkg/mxutil"
)

// joinChunks concatenates a slice of byte slices in one allocation.
// Uses byte-by-byte copy because s[off:] creates a copy under
// subslice-copies semantics, so copy(s[off:], src) writes to a
// temporary instead of s.
func joinChunks(chunks [][]byte) (r []byte) {
	n := int32(0)
	for _, c := range chunks {
		n += len(c)
	}
	out := []byte{:n}
	off := int32(0)
	for _, c := range chunks {
		for j := int32(0); j < len(c); j++ {
			out[off+j] = c[j]
		}
		off += len(c)
	}
	return out
}

// commaSep interleaves ", " between segments.
func commaSep(segs []string) (r []string) {
	var out []string
	for i, s := range segs {
		if i > 0 {
			push(out, ", ")
		}
		push(out, s)
	}
	return out
}

// joinStrs concatenates string segments in one allocation.
// Replaces Schlemiel x = x | y loops.
func joinStrs(segs []string) (r string) {
	n := int32(0)
	for _, s := range segs {
		n += len(s)
	}
	out := []byte{:n}
	off := int32(0)
	for _, s := range segs {
		for j := int32(0); j < len(s); j++ {
			out[off+j] = s[j]
		}
		off += len(s)
	}
	return string(out)
}

// srcTrimStart returns index of first non-whitespace byte in src[from:to].
func srcTrimStart(src []byte, from, to int32) (s int32) {
	s = from
	for s < to && (src[s] == ' ' || src[s] == '\t' || src[s] == '\r') {
		s++
	}
	return s
}

// srcTrimEnd returns index one past last non-whitespace byte in src[from:to].
func srcTrimEnd(src []byte, from, to int32) (e int32) {
	e = to
	for e > from && (src[e-1] == ' ' || src[e-1] == '\t' || src[e-1] == '\r') {
		e--
	}
	return e
}

// srcHasPrefix checks whether src[start:end] begins with prefix.
func srcHasPrefix(src []byte, start, end int32, prefix []byte) (ok bool) {
	plen := int32(len(prefix))
	if end-start < plen {
		return false
	}
	for i := int32(0); i < plen; i++ {
		if src[start+i] != prefix[i] {
			return false
		}
	}
	return true
}

// srcFindByte finds byte c in src[from:to], returns index or -1.
func srcFindByte(src []byte, from, to int32, c byte) (idx int32) {
	for i := from; i < to; i++ {
		if src[i] == c {
			return i
		}
	}
	return -1
}

func scanEmbedPragmas(src []byte, pkgDir string) {
	cctx.embedVars = map[string]*embedInfo{}
	embedPrefix := []byte("//moxie:embed_le ")
	embedPrefix2 := []byte("//:embed ")
	varPrefix := []byte("var ")
	commentPrefix := []byte("//")
	pendingPath := ""
	slen := int32(len(src))
	i := int32(0)
	for i < slen {
		// Find line boundaries without subslicing.
		le := srcFindByte(src, i, slen, '\n')
		if le < 0 {
			le = slen
		}
		ts := srcTrimStart(src, i, le)
		te := srcTrimEnd(src, i, le)
		tlen := te - ts
		if len(pendingPath) > 0 {
			if tlen == 0 || srcHasPrefix(src, ts, te, commentPrefix) {
				if le < slen {
					i = le + 1
				} else {
					i = slen
				}
				continue
			}
			if srcHasPrefix(src, ts, te, varPrefix) {
				rs := ts + 4
				sp := srcFindByte(src, rs, te, ' ')
				if sp > rs {
					varName := string(src[rs:sp])
					vs := srcTrimStart(src, sp+1, te)
					ve := srcTrimEnd(src, vs, te)
					typeStr := string(src[vs:ve])
					cctx.embedVars[varName] = &embedInfo{path: pendingPath, elemType: typeStr}
				}
			}
			pendingPath = ""
		} else if srcHasPrefix(src, ts, te, embedPrefix) {
			ps := srcTrimStart(src, ts+int32(len(embedPrefix)), te)
			pe := srcTrimEnd(src, ps, te)
			if pe > ps {
				pendingPath = string(src[ps:pe])
				if pkgDir != "" {
					pendingPath = mxutil.JoinPath(pkgDir, pendingPath)
				}
			}
		} else if srcHasPrefix(src, ts, te, embedPrefix2) {
			ps := srcTrimStart(src, ts+int32(len(embedPrefix2)), te)
			pe := srcTrimEnd(src, ps, te)
			if pe > ps {
				pendingPath = string(src[ps:pe])
				if pkgDir != "" {
					pendingPath = mxutil.JoinPath(pkgDir, pendingPath)
				}
			}
		}
		if le < slen {
			i = le + 1
		} else {
			i = slen
		}
	}
}

func scanExportPragmas(src []byte) (m map[string]string) {
	result := map[string]string{}
	exportPrefix := []byte("//export ")
	funcPrefix := []byte("func ")
	commentPrefix := []byte("//")
	pendingExport := ""
	slen := int32(len(src))
	i := int32(0)
	for i < slen {
		le := srcFindByte(src, i, slen, '\n')
		if le < 0 {
			le = slen
		}
		ts := srcTrimStart(src, i, le)
		te := srcTrimEnd(src, i, le)
		tlen := te - ts
		if len(pendingExport) > 0 {
			if tlen == 0 || srcHasPrefix(src, ts, te, commentPrefix) {
				if le < slen { i = le + 1 } else { i = slen }
				continue
			}
			if srcHasPrefix(src, ts, te, funcPrefix) {
				rs := ts + 5
				paren := srcFindByte(src, rs, te, '(')
				if paren > rs {
					ns := srcTrimStart(src, rs, paren)
					ne := srcTrimEnd(src, ns, paren)
					if ne > ns {
						result[string(src[ns:ne])] = pendingExport
					}
				}
			}
			pendingExport = ""
		} else if srcHasPrefix(src, ts, te, exportPrefix) {
			ps := srcTrimStart(src, ts+9, te)
			pe := srcTrimEnd(src, ps, te)
			if pe > ps {
				pendingExport = string(src[ps:pe])
			}
		}
		if le < slen { i = le + 1 } else { i = slen }
	}
	return result
}

func scanLinknameMap(src []byte) (m map[string]string) {
	result := map[string]string{}
	prefix := []byte("//:linkname ")
	slen := int32(len(src))
	i := int32(0)
	for i < slen {
		le := srcFindByte(src, i, slen, '\n')
		if le < 0 {
			le = slen
		}
		ts := srcTrimStart(src, i, le)
		te := srcTrimEnd(src, i, le)
		if srcHasPrefix(src, ts, te, prefix) {
			rs := srcTrimStart(src, ts+12, te)
			re := srcTrimEnd(src, rs, te)
			if re > rs {
				// Find space between localName and targetName.
				sp := srcFindByte(src, rs, re, ' ')
				if sp < 0 {
					sp = srcFindByte(src, rs, re, '\t')
				}
				if sp > rs {
					localName := string(src[rs:sp])
					ts2 := srcTrimStart(src, sp+1, re)
					if ts2 < re {
						targetName := string(src[ts2:re])
						result[localName] = targetName
					}
				}
			}
		}
		if le < slen { i = le + 1 } else { i = slen }
	}
	return result
}


// wasmImportMap maps IR symbol (e.g. "time.timezoneOffsetMinutes") -> [module, import_name].
// Populated from //:wasmimport pragmas.

// scanWasmImportMap scans source for //:wasmimport <module> <name> pragmas.
// The pragma applies to the next function declaration (body-less).
// pkgPath is the package import path used to construct full IR symbols.
func scanWasmImportMap(src []byte, pkgPath string) (m map[string][2]string) {
	if !mxutil.HasPrefix(cctx.compileTargetTriple, "wasm32") {
		return nil
	}
	result := map[string][2]string{}
	wasmPrefix := []byte("//:wasmimport ")
	funcPrefix := []byte("func ")
	slen := int32(len(src))
	i := int32(0)
	for i < slen {
		le := srcFindByte(src, i, slen, '\n')
		if le < 0 {
			le = slen
		}
		ts := srcTrimStart(src, i, le)
		te := srcTrimEnd(src, i, le)
		if !srcHasPrefix(src, ts, te, wasmPrefix) {
			if le < slen { i = le + 1 } else { i = slen }
			continue
		}
		// Parse module and import name from pragma line.
		rs := ts + int32(len(wasmPrefix))
		sp := srcFindByte(src, rs, te, ' ')
		if sp < 0 {
			sp = srcFindByte(src, rs, te, '\t')
		}
		if sp < 0 {
			if le < slen { i = le + 1 } else { i = slen }
			continue
		}
		mod := string(src[rs:sp])
		ns := srcTrimStart(src, sp+1, te)
		ne := srcTrimEnd(src, ns, te)
		impName := string(src[ns:ne])
		// Find next func declaration.
		j := le
		if j < slen { j++ }
		for j < slen {
			fle := srcFindByte(src, j, slen, '\n')
			if fle < 0 {
				fle = slen
			}
			fts := srcTrimStart(src, j, fle)
			fte := srcTrimEnd(src, j, fle)
			if !srcHasPrefix(src, fts, fte, funcPrefix) {
				if fle < slen { j = fle + 1 } else { j = slen }
				continue
			}
			// Extract function name from "func NAME(" or "func (recv) NAME("
			rs2 := fts + 5
			if rs2 < fte && src[rs2] == '(' {
				depth := int32(1)
				k := rs2 + 1
				for k < fte && depth > 0 {
					if src[k] == '(' { depth++ }
					if src[k] == ')' { depth-- }
					k++
				}
				rs2 = srcTrimStart(src, k, fte)
			}
			paren := srcFindByte(src, rs2, fte, '(')
			if paren < 0 {
				break
			}
			fns := srcTrimStart(src, rs2, paren)
			fne := srcTrimEnd(src, fns, paren)
			if fne > fns {
				funcName := string(src[fns:fne])
				sym := pkgPath | "." | funcName
				result[sym] = [2]string{mod, impName}
			}
			break
		}
		if le < slen { i = le + 1 } else { i = slen }
	}
	if len(result) == 0 {
		return nil
	}
	return result
}

func stripDuplicatePackageClauses(src []byte) (buf []byte) {
	found := false
	var out []byte
	i := 0
	for i < len(src) {
		nlIdx := lIndexByte(src[i:], '\n')
		var line []byte
		var lineEnd int32
		if nlIdx < 0 {
			line = src[i:]
			lineEnd = len(src)
		} else {
			line = src[i : i+nlIdx]
			lineEnd = i + nlIdx + 1
		}
		trimmed := lTrimSpace(line)
		if lHasPrefix(trimmed, "package ") {
			if found {
				if out == nil {
					out = []byte{:len(src)}
					copy(out, src)
				}
				for k := i; k < i+len(line); k++ {
					out[k] = ' '
				}
				i = lineEnd
				continue
			}
			found = true
		}
		i = lineEnd
	}
	if out == nil {
		return src
	}
	return out
}

func stripDuplicatePackageClauses2(src []byte) (buf []byte) {
	// Fast path: count package clauses. If only one, return src unchanged.
	count := int32(0)
	for i := int32(0); i < len(src); i++ {
		if i == 0 || src[i-1] == '\n' {
			j := i
			for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
				j++
			}
			if len(src)-j >= 8 && src[j] == 'p' && src[j+1] == 'a' && src[j+2] == 'c' && src[j+3] == 'k' && src[j+4] == 'a' && src[j+5] == 'g' && src[j+6] == 'e' && src[j+7] == ' ' {
				count++
			}
		}
	}
	if count <= 1 {
		return src
	}
	// Multiple package clauses: copy src, blank out duplicates in-place.
	out := []byte{:len(src)}
	copy(out, src)
	found := false
	for i := int32(0); i < len(out); i++ {
		if i == 0 || out[i-1] == '\n' {
			ls := i
			for ls < len(out) && (out[ls] == ' ' || out[ls] == '\t') {
				ls++
			}
			if len(out)-ls >= 8 && out[ls] == 'p' && out[ls+1] == 'a' && out[ls+2] == 'c' && out[ls+3] == 'k' && out[ls+4] == 'a' && out[ls+5] == 'g' && out[ls+6] == 'e' && out[ls+7] == ' ' {
				if found {
					// Blank out this line.
					for k := i; k < len(out) && out[k] != '\n'; k++ {
						out[k] = ' '
					}
				}
				found = true
			}
		}
	}
	return out
}

func rewriteSliceMakeLiterals(src []byte) (buf []byte) {
	var chunks [][]byte
	i := 0
	for i < len(src) {
		start := lIndex(src[i:], []byte("{:"))
		if start < 0 {
			push(chunks, src[i:])
			break
		}
		start = start + i
		lbrack := findSliceTypeStart(src, start)
		if lbrack < 0 {
			push(chunks, src[i:start+2])
			i = start + 2
			continue
		}
		closeIdx := findMatchingBrace(src, start)
		if closeIdx < 0 {
			push(chunks, src[i:start+2])
			i = start + 2
			continue
		}
		inner := src[start+2 : closeIdx]
		colonIdx := lIndexByte(inner, ':')
		typeText := src[lbrack:start]
		push(chunks, src[i:lbrack])
		if colonIdx < 0 {
			push(chunks, []byte("__slicealloc("))
			push(chunks, typeText)
			push(chunks, []byte(", "))
			push(chunks, lTrimSpace(inner))
			push(chunks, []byte(")"))
		} else {
			lenExpr := lTrimSpace(inner[:colonIdx])
			capExpr := lTrimSpace(inner[colonIdx+1:])
			push(chunks, []byte("__slicealloc("))
			push(chunks, typeText)
			push(chunks, []byte(", "))
			push(chunks, lenExpr)
			push(chunks, []byte(", "))
			push(chunks, capExpr)
			push(chunks, []byte(")"))
		}
		i = closeIdx + 1
	}
	if len(chunks) == 0 {
		return src
	}
	return joinChunks(chunks)
}

func findSliceTypeStart(src []byte, braceIdx int32) (n int32) {
	j := braceIdx - 1
	for j >= 0 && (src[j] == ' ' || src[j] == '\t' || src[j] == '\n') {
		j--
	}
	if j < 0 {
		return -1
	}
	depth := 0
	parenDepth := 0
	candidate := -1
	for j >= 0 {
		ch := src[j]
		if ch == ')' {
			parenDepth++
		} else if ch == '(' {
			if parenDepth > 0 {
				parenDepth--
			} else {
				if candidate >= 0 {
					return candidate
				}
				return -1
			}
		} else if parenDepth > 0 {
			j--
			continue
		}
		if ch == ']' {
			depth++
		} else if ch == '[' {
			depth--
			if depth == 0 {
				candidate = j
			}
		} else if depth == 0 && parenDepth == 0 {
			if candidate >= 0 {
				return candidate
			}
			if ch == ' ' || ch == '\t' || ch == '\n' || ch == '*' || ch == '(' || ch == ')' {
				j--
				continue
			}
			if (ch >= 'a' && ch <= 'z') || (ch >= 'A' && ch <= 'Z') || (ch >= '0' && ch <= '9') || ch == '_' || ch == '.' {
				j--
				continue
			}
			return -1
		}
		j--
	}
	if candidate >= 0 {
		return candidate
	}
	return -1
}

func findMatchingBrace(src []byte, openIdx int32) (n int32) {
	depth := 1
	for i := openIdx + 1; i < len(src); i++ {
		if src[i] == '{' {
			depth++
		} else if src[i] == '}' {
			depth--
			if depth == 0 {
				return i
			}
		}
	}
	return -1
}

func rewriteChanMakeLiterals(src []byte) (buf []byte) {
	chanKw := []byte("chan ")
	var chunks [][]byte
	started := false
	i := 0
	for i < len(src) {
		idx := lIndex(src[i:], chanKw)
		if idx < 0 {
			if started {
				push(chunks, src[i:])
			}
			break
		}
		idx = idx + i
		j := idx + 5
		for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
			j++
		}
		if j >= len(src) {
			if started {
				push(chunks, src[i:])
			}
			break
		}
		if src[j] == '<' && j+1 < len(src) && src[j+1] == '-' {
			if started {
				push(chunks, src[i:j+2])
			}
			i = j + 2
			continue
		}
		for j < len(src) && (src[j] != '{' && src[j] != '\n' && src[j] != ';' && src[j] != '(' && src[j] != ')') {
			if src[j] == ' ' || src[j] == '\t' {
				break
			}
			j++
		}
		if j >= len(src) || src[j] != '{' {
			if !started {
				i = idx + 4
			} else {
				push(chunks, src[i:idx+4])
				i = idx + 4
			}
			continue
		}
		braceOpen := j
		endsStruct := braceOpen >= 6 && string(src[braceOpen-6:braceOpen]) == "struct"
		endsIface := braceOpen >= 9 && string(src[braceOpen-9:braceOpen]) == "interface"
		if endsStruct || endsIface {
			typeClose := findMatchingBrace(src, braceOpen)
			if typeClose < 0 {
				if !started {
					i = idx + 4
				} else {
					push(chunks, src[i:idx+4])
					i = idx + 4
				}
				continue
			}
			k := typeClose + 1
			if k >= len(src) || src[k] != '{' {
				if !started {
					i = typeClose + 1
				} else {
					push(chunks, src[i:typeClose+1])
					i = typeClose + 1
				}
				continue
			}
			braceOpen = k
		}
		closeIdx := findMatchingBrace(src, braceOpen)
		if closeIdx < 0 {
			if !started {
				i = idx + 4
			} else {
				push(chunks, src[i:idx+4])
				i = idx + 4
			}
			continue
		}
		inner := lTrimSpace(src[braceOpen+1 : closeIdx])
		chanType := src[idx : braceOpen]
		for len(chanType) > 0 && (chanType[len(chanType)-1] == ' ' || chanType[len(chanType)-1] == '\t') {
			chanType = chanType[:len(chanType)-1]
		}
		if !started {
			push(chunks, src[:idx])
			started = true
		} else {
			push(chunks, src[i:idx])
		}
		if len(inner) == 0 {
			push(chunks, []byte("__slicealloc("))
			push(chunks, chanType)
			push(chunks, []byte(")"))
		} else {
			push(chunks, []byte("__slicealloc("))
			push(chunks, chanType)
			push(chunks, []byte(", "))
			push(chunks, inner)
			push(chunks, []byte(")"))
		}
		i = closeIdx + 1
	}
	if !started {
		return src
	}
	return joinChunks(chunks)
}