package iskra

import "bytes"

type IRSignature struct {
	ReturnType string
	Params     []IRParam
	FuncName   string
}

type IRParam struct {
	Type string
	Name string
}

func ParseIRSignature(defLine []byte) IRSignature {
	sig := IRSignature{}
	// "define hidden i1 @"pkg.Func"(i8 %c, ptr %context) ..."
	rest := defLine

	// Skip "define" and attributes
	for len(rest) > 0 && rest[0] != '@' {
		// find return type: the token before @
		idx := bytes.IndexByte(rest, '@')
		if idx < 0 {
			return sig
		}
		before := bytes.TrimSpace(rest[:idx])
		// return type is last space-separated token before @
		sp := bytes.LastIndexByte(before, ' ')
		if sp >= 0 {
			sig.ReturnType = string(before[sp+1:])
		} else {
			sig.ReturnType = string(before)
		}
		rest = rest[idx:]
		break
	}

	// Extract function name
	if len(rest) > 0 && rest[0] == '@' {
		rest = rest[1:]
		if len(rest) > 0 && rest[0] == '"' {
			rest = rest[1:]
			end := bytes.IndexByte(rest, '"')
			if end >= 0 {
				sig.FuncName = string(rest[:end])
				rest = rest[end+1:]
			}
		} else {
			end := bytes.IndexByte(rest, '(')
			if end >= 0 {
				sig.FuncName = string(rest[:end])
				rest = rest[end:]
			}
		}
	}

	// Parse parameters
	pOpen := bytes.IndexByte(rest, '(')
	if pOpen < 0 {
		return sig
	}
	rest = rest[pOpen+1:]
	pClose := bytes.IndexByte(rest, ')')
	if pClose < 0 {
		return sig
	}
	paramStr := rest[:pClose]

	parts := bytes.Split(paramStr, []byte(","))
	for _, p := range parts {
		p = bytes.TrimSpace(p)
		if len(p) == 0 {
			continue
		}
		param := IRParam{}
		// Strip attributes like "dereferenceable_or_null(N)"
		p = stripParamAttrs(p)
		sp := bytes.LastIndexByte(p, ' ')
		if sp >= 0 {
			param.Type = string(bytes.TrimSpace(p[:sp]))
			param.Name = string(bytes.TrimSpace(p[sp+1:]))
		} else {
			param.Type = string(p)
		}
		sig.Params = append(sig.Params, param)
	}
	return sig
}

func stripParamAttrs(p []byte) []byte {
	// Remove "dereferenceable_or_null(N)" and "dereferenceable(N)"
	for {
		idx := bytes.Index(p, []byte("dereferenceable"))
		if idx < 0 {
			break
		}
		end := idx + 15
		if end < len(p) && p[end] == '_' {
			// dereferenceable_or_null
			nEnd := bytes.Index(p[idx:], []byte(")"))
			if nEnd >= 0 {
				end = idx + nEnd + 1
			}
		} else if end < len(p) && p[end] == '(' {
			nEnd := bytes.IndexByte(p[end:], ')')
			if nEnd >= 0 {
				end = end + nEnd + 1
			}
		}
		var rebuilt []byte
		rebuilt = append(rebuilt, p[:idx]...)
		rebuilt = append(rebuilt, p[end:]...)
		p = rebuilt
	}
	// Collapse double spaces
	for bytes.Contains(p, []byte("  ")) {
		p = bytes.Replace(p, []byte("  "), []byte(" "), -1)
	}
	return bytes.TrimSpace(p)
}

type VerifyCase struct {
	FuncName  string
	Sig       IRSignature
	Testable  bool
	TestKind  string // "byte-exhaustive", "int-boundary", "void-null", "untestable"
	LatticeIR []byte
	RefIR     []byte
}

func ClassifyTestability(sig IRSignature) (string, bool) {
	if len(sig.Params) == 0 {
		return "untestable", false
	}
	// Last param is always ptr %context - ignore it
	realParams := sig.Params
	if len(realParams) > 0 && realParams[len(realParams)-1].Name == "%context" {
		realParams = realParams[:len(realParams)-1]
	}

	if len(realParams) == 0 {
		if sig.ReturnType == "i1" || sig.ReturnType == "i32" {
			return "const-null", true
		}
		if sig.ReturnType == "void" {
			return "void-null", true
		}
		return "untestable", false
	}

	if len(realParams) == 1 {
		t := realParams[0].Type
		if t == "i8" && (sig.ReturnType == "i1" || sig.ReturnType == "i32") {
			return "byte-exhaustive", true
		}
		if t == "i32" && (sig.ReturnType == "i1" || sig.ReturnType == "i32") {
			return "int-boundary", true
		}
	}

	// Multi-scalar params
	allScalar := true
	for _, p := range realParams {
		if p.Type != "i8" && p.Type != "i16" && p.Type != "i32" && p.Type != "i64" && p.Type != "i1" {
			allScalar = false
			break
		}
	}
	if allScalar && (sig.ReturnType == "i1" || sig.ReturnType == "i32" || sig.ReturnType == "i64" || sig.ReturnType == "void") {
		return "scalar-boundary", true
	}

	return "untestable", false
}

func GenerateTestMain(sig IRSignature, testKind string, needDecl bool) []byte {
	var b []byte
	b = append(b, "declare i32 @printf(ptr, ...)\n"...)
	b = append(b, "@fmt_d = private unnamed_addr constant [4 x i8] c\"%d\\0A\\00\"\n"...)
	b = append(b, "@fmt_ld = private unnamed_addr constant [5 x i8] c\"%ld\\0A\\00\"\n"...)

	if needDecl {
		b = append(b, "declare " | sig.ReturnType | " @\"" | sig.FuncName | "\"("...)
		for i, p := range sig.Params {
			if i > 0 {
				b = append(b, ", "...)
			}
			b = append(b, p.Type...)
		}
		b = append(b, ")\n\n"...)
	}

	b = append(b, "define i32 @main() {\n"...)
	b = append(b, "entry:\n"...)

	switch testKind {
	case "byte-exhaustive":
		b = append(b, "  br label %loop\n"...)
		b = append(b, "loop:\n"...)
		b = append(b, "  %i = phi i32 [ 0, %entry ], [ %next, %loop ]\n"...)
		b = append(b, "  %c = trunc i32 %i to i8\n"...)
		b = append(b, "  %r = call " | sig.ReturnType | " @\"" | sig.FuncName | "\"(i8 %c, ptr null)\n"...)
		if sig.ReturnType == "i1" {
			b = append(b, "  %r32 = zext i1 %r to i32\n"...)
			b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %r32)\n"...)
		} else {
			b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %r)\n"...)
		}
		b = append(b, "  %next = add i32 %i, 1\n"...)
		b = append(b, "  %done = icmp eq i32 %next, 256\n"...)
		b = append(b, "  br i1 %done, label %exit, label %loop\n"...)
		b = append(b, "exit:\n"...)
		b = append(b, "  ret i32 0\n"...)

	case "int-boundary":
		vals := []string{"0", "1", "-1", "127", "-128", "255", "2147483647", "-2147483648", "42", "100", "1000"}
		for vi, v := range vals {
			vn := "v" | intToStr(vi)
			b = append(b, "  %" | vn | " = call " | sig.ReturnType | " @\"" | sig.FuncName | "\"(i32 " | v | ", ptr null)\n"...)
			if sig.ReturnType == "i1" {
				b = append(b, "  %" | vn | "e = zext i1 %" | vn | " to i32\n"...)
				b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %" | vn | "e)\n"...)
			} else {
				b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %" | vn | ")\n"...)
			}
		}
		b = append(b, "  ret i32 0\n"...)

	case "const-null":
		b = append(b, "  %r = call " | sig.ReturnType | " @\"" | sig.FuncName | "\"(ptr null)\n"...)
		if sig.ReturnType == "i1" {
			b = append(b, "  %r32 = zext i1 %r to i32\n"...)
			b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %r32)\n"...)
		} else {
			b = append(b, "  call i32 (ptr, ...) @printf(ptr @fmt_d, i32 %r)\n"...)
		}
		b = append(b, "  ret i32 0\n"...)

	case "void-null":
		b = append(b, "  call void @\"" | sig.FuncName | "\"(ptr null)\n"...)
		b = append(b, "  ret i32 0\n"...)

	default:
		b = append(b, "  ret i32 0\n"...)
	}

	b = append(b, "}\n"...)
	return b
}

func ReplaceFunctionInModule(module []byte, funcName string, newBody []byte) []byte {
	// Find "define ... @"funcName"(...) ... {"
	searchStr := "@\"" | funcName | "\""
	defIdx := bytes.Index(module, []byte("define"))
	for defIdx >= 0 {
		nameIdx := bytes.Index(module[defIdx:], []byte(searchStr))
		if nameIdx < 0 {
			break
		}
		nameIdx += defIdx
		// Verify this is within a define line (find the newline before)
		lineStart := bytes.LastIndexByte(module[:nameIdx], '\n')
		if lineStart < 0 {
			lineStart = 0
		} else {
			lineStart++
		}
		line := module[lineStart:]
		if !bytes.HasPrefix(bytes.TrimSpace(line), []byte("define ")) {
			defIdx = nameIdx + len(searchStr)
			continue
		}

		// Find the opening {
		braceIdx := bytes.IndexByte(module[nameIdx:], '{')
		if braceIdx < 0 {
			break
		}
		funcStart := lineStart
		bodyStart := nameIdx + braceIdx

		// Find the closing } - match braces
		depth := 1
		pos := bodyStart + 1
		for pos < len(module) && depth > 0 {
			if module[pos] == '{' {
				depth++
			} else if module[pos] == '}' {
				depth--
			}
			pos++
		}
		funcEnd := pos

		// Build replacement
		var result []byte
		result = append(result, module[:funcStart]...)
		result = append(result, newBody...)
		result = append(result, '\n')
		result = append(result, module[funcEnd:]...)
		return result
	}
	return module
}

func ExtractFunctionFromModule(module []byte, funcName string) []byte {
	searchStr := "@\"" | funcName | "\""
	defIdx := 0
	for defIdx < len(module) {
		nextDef := bytes.Index(module[defIdx:], []byte("define"))
		if nextDef < 0 {
			break
		}
		defIdx += nextDef
		nameIdx := bytes.Index(module[defIdx:], []byte(searchStr))
		if nameIdx < 0 || nameIdx > 200 {
			defIdx += 7
			continue
		}
		nameIdx += defIdx

		lineStart := bytes.LastIndexByte(module[:nameIdx], '\n')
		if lineStart < 0 {
			lineStart = 0
		} else {
			lineStart++
		}

		braceIdx := bytes.IndexByte(module[nameIdx:], '{')
		if braceIdx < 0 {
			break
		}
		bodyStart := nameIdx + braceIdx
		depth := 1
		pos := bodyStart + 1
		for pos < len(module) && depth > 0 {
			if module[pos] == '{' {
				depth++
			} else if module[pos] == '}' {
				depth--
			}
			pos++
		}
		return module[lineStart:pos]
	}
	return nil
}