package rewrite

import "bytes"

func StripDuplicatePackageClauses(src []byte) []byte {
	found := false
	var out []byte
	i := 0
	for i < len(src) {
		nlIdx := bytes.IndexByte(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 := bytes.TrimSpace(line)
		if bytes.HasPrefix(trimmed, "package ") {
			if found {
				if out == nil {
					out = []byte{:0:len(src)}
					out = append(out, src[:i]...)
				}
				for k := 0; k < len(line); k++ {
					out = append(out, ' ')
				}
				if nlIdx >= 0 {
					out = append(out, '\n')
				}
				i = lineEnd
				continue
			}
			found = true
		}
		if out != nil {
			out = append(out, src[i:lineEnd]...)
		}
		i = lineEnd
	}
	if out == nil {
		return src
	}
	return out
}

func RewriteSliceMakeLiterals(src []byte) []byte {
	var out []byte
	i := 0
	for i < len(src) {
		start := bytes.Index(src[i:], []byte("{:"))
		if start < 0 {
			out = append(out, src[i:]...)
			break
		}
		start = start + i
		lbrack := findSliceTypeStart(src, start)
		if lbrack < 0 {
			out = append(out, src[i:start+2]...)
			i = start + 2
			continue
		}
		close := findMatchingBrace(src, start)
		if close < 0 {
			out = append(out, src[i:start+2]...)
			i = start + 2
			continue
		}
		inner := src[start+2 : close]
		colonIdx := bytes.IndexByte(inner, ':')
		typeText := src[lbrack:start]
		out = append(out, src[i:lbrack]...)
		if colonIdx < 0 {
			out = append(out, "make("...)
			out = append(out, typeText...)
			out = append(out, ", "...)
			out = append(out, bytes.TrimSpace(inner)...)
			out = append(out, ')')
		} else {
			lenExpr := bytes.TrimSpace(inner[:colonIdx])
			capExpr := bytes.TrimSpace(inner[colonIdx+1:])
			out = append(out, "make("...)
			out = append(out, typeText...)
			out = append(out, ", "...)
			out = append(out, lenExpr...)
			out = append(out, ", "...)
			out = append(out, capExpr...)
			out = append(out, ')')
		}
		i = close + 1
	}
	if out == nil {
		return src
	}
	return out
}

func findSliceTypeStart(src []byte, braceIdx int32) 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 {
				if candidate >= 0 {
					return candidate
				}
				return -1
			}
			parenDepth--
		} else if parenDepth > 0 {
			j--
			continue
		}
		if ch == ']' {
			depth++
		} else if ch == '[' {
			depth--
			if depth == 0 {
				candidate = j
			}
		} else if ch == '{' || ch == ';' {
			if candidate >= 0 {
				return candidate
			}
			return -1
		} 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) 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) []byte {
	chanKw := []byte("chan ")
	var out []byte
	i := 0
	for i < len(src) {
		idx := bytes.Index(src[i:], chanKw)
		if idx < 0 {
			if out != nil {
				out = append(out, src[i:]...)
			}
			break
		}
		idx = idx + i
		j := idx + 5
		for j < len(src) && (src[j] == ' ' || src[j] == '\t') {
			j++
		}
		if j >= len(src) {
			if out != nil {
				out = append(out, src[i:]...)
			}
			break
		}
		if src[j] == '<' && j+1 < len(src) && src[j+1] == '-' {
			if out != nil {
				out = append(out, 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 out == nil {
				i = idx + 4
			} else {
				out = append(out, src[i:idx+4]...)
				i = idx + 4
			}
			continue
		}
		braceOpen := j
		endsStruct := braceOpen >= 6 && string(src[braceOpen-6:braceOpen]) == "struct"
		if !endsStruct && braceOpen >= 7 {
			k := braceOpen - 1
			for k > idx && (src[k] == ' ' || src[k] == '\t') {
				k--
			}
			if k >= 5 && string(src[k-5:k+1]) == "struct" {
				endsStruct = true
			}
		}
		if endsStruct {
			structClose := findMatchingBrace(src, braceOpen)
			if structClose < 0 {
				if out == nil {
					i = idx + 4
				} else {
					out = append(out, src[i:idx+4]...)
					i = idx + 4
				}
				continue
			}
			nj := structClose + 1
			if nj >= len(src) || src[nj] != '{' {
				if out == nil {
					i = structClose + 1
				} else {
					out = append(out, src[i:structClose+1]...)
					i = structClose + 1
				}
				continue
			}
			braceOpen = nj
		}
		close := findMatchingBrace(src, braceOpen)
		if close < 0 {
			if out == nil {
				i = idx + 4
			} else {
				out = append(out, src[i:idx+4]...)
				i = idx + 4
			}
			continue
		}
		inner := bytes.TrimSpace(src[braceOpen+1 : close])
		chanType := src[idx : braceOpen]
		for len(chanType) > 0 && (chanType[len(chanType)-1] == ' ' || chanType[len(chanType)-1] == '\t') {
			chanType = chanType[:len(chanType)-1]
		}
		if out == nil {
			out = []byte{:0:len(src)}
			out = append(out, src[:idx]...)
		} else {
			out = append(out, src[i:idx]...)
		}
		if len(inner) == 0 {
			out = append(out, "make("...)
			out = append(out, chanType...)
			out = append(out, ')')
		} else {
			out = append(out, "make("...)
			out = append(out, chanType...)
			out = append(out, ", "...)
			out = append(out, inner...)
			out = append(out, ')')
		}
		i = close + 1
	}
	if out == nil {
		return src
	}
	return out
}