"""Classify strict COFF codegen differences without granting match credit.

The classifier is deliberately fail-closed.  It maps a small set of measured
MSVC 7.0 code-generation signatures to the production recipes documented by
the XDK 3911 codegen atlas.  Anything that does not satisfy a complete rule is
``UNKNOWN``.  A result is diagnostic evidence only: this tool never changes a
matching ledger or decides that a function is exact.
"""

import argparse
import itertools
import json
import re
import shutil
import subprocess
from pathlib import Path

from .coff_compare import (
	CoffError,
	load,
	relocation_infos_equal,
	section_info,
	section_infos_equal,
)


VERSION = 1
GENERAL_REGISTERS = ("eax", "ebx", "ecx", "edx", "esi", "edi")
REGISTER_ALIASES = {
	"eax": ("eax", "ax", "al", "ah"),
	"ebx": ("ebx", "bx", "bl", "bh"),
	"ecx": ("ecx", "cx", "cl", "ch"),
	"edx": ("edx", "dx", "dl", "dh"),
	"esi": ("esi", "si"),
	"edi": ("edi", "di"),
	"ebp": ("ebp", "bp"),
	"esp": ("esp", "sp"),
}
ALIAS_TO_REGISTER = {
	alias: register
	for register, aliases in REGISTER_ALIASES.items()
	for alias in aliases
}
INSTRUCTION_RE = re.compile(
	r"^\s*([0-9a-fA-F]+):\s+((?:[0-9a-fA-F]{2}\s+)+)"
	r"([A-Za-z][A-Za-z0-9_.]*)\s*(.*?)\s*$")
REGISTER_RE = re.compile(
	r"%(eax|ax|al|ah|ebx|bx|bl|bh|ecx|cx|cl|ch|edx|dx|dl|dh|"
	r"esi|si|edi|di|ebp|bp|esp|sp)\b")


class ClassifierError(RuntimeError):
	pass


def parse_disassembly(text):
	"""Parse llvm-objdump AT&T instruction lines into stable evidence."""
	instructions = []
	for line in text.splitlines():
		match = INSTRUCTION_RE.match(line)
		if not match:
			continue
		instructions.append({
			"address": int(match.group(1), 16),
			"bytes": "".join(match.group(2).split()).lower(),
			"mnemonic": match.group(3).lower(),
			"operands": _normalize_operands(match.group(4)),
		})
	return instructions


def _normalize_operands(operands):
	return re.sub(r"\s+", " ", operands.strip().lower())


def _instruction_key(instruction):
	return instruction["mnemonic"], _normalize_operands(instruction["operands"])


def _split_operands(operands):
	"""Split AT&T operands without treating SIB commas as separators."""
	parts = []
	start = 0
	depth = 0
	for index, character in enumerate(operands):
		if character == "(":
			depth += 1
		elif character == ")":
			depth = max(0, depth - 1)
		elif character == "," and depth == 0:
			parts.append(operands[start:index].strip())
			start = index + 1
	parts.append(operands[start:].strip())
	return parts


def _relocation_identities_equal(left, right):
	"""Compare ordered relocation semantics while allowing address drift."""
	if len(left) != len(right):
		return False
	left_without_addresses = [dict(relocation, address=0) for relocation in left]
	right_without_addresses = [dict(relocation, address=0) for relocation in right]
	return relocation_infos_equal(left_without_addresses, right_without_addresses)


def _metadata(target, candidate):
	target_relocations = target.get("relocations", [])
	candidate_relocations = candidate.get("relocations", [])
	return {
		"size_equal": target.get("size") == candidate.get("size"),
		"relocation_count_equal": (
			target.get("relocation_count") == candidate.get("relocation_count")),
		"relocation_identities_equal": _relocation_identities_equal(
			target_relocations, candidate_relocations),
		"relocations_exact": relocation_infos_equal(
			target_relocations, candidate_relocations),
	}


def _result(kind, blocker_class, recipe, confidence, evidence, metadata,
			stop_rule, preconditions=None):
	return {
		"version": VERSION,
		"kind": kind,
		"class": blocker_class,
		"recipe": recipe,
		"confidence": confidence,
		"evidence": evidence,
		"strict_metadata": metadata,
		"preconditions": preconditions or [],
		"stop_rule": stop_rule,
		"grants_match_credit": False,
	}


def classify(target, candidate):
	"""Return one conservative codegen classification for two evidence dicts."""
	metadata = _metadata(target, candidate)
	target_instructions = target.get("instructions", [])
	candidate_instructions = candidate.get("instructions", [])
	if not target_instructions or not candidate_instructions:
		return _unknown(metadata, "both disassemblies are required")

	if target.get("normalized_sha256") == candidate.get("normalized_sha256") \
			and metadata["size_equal"] and metadata["relocations_exact"]:
		return _result(
			"REFUSE", "no-codegen-blocker", None, "PROVEN",
			["strict measurements are already equal"], metadata,
			"Use the ordinary strict comparator; this classifier does not grant credit.")

	classification = _classify_return_shape(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	classification = _classify_private_abi(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	classification = _classify_alias_shape(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	classification = _classify_x87_liveness(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	classification = _classify_register_permutation(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	classification = _classify_scheduling_tie(
		target_instructions, candidate_instructions, metadata)
	if classification:
		return classification
	return _unknown(metadata, "no proven recipe or strict STOP signature matched")


def _unknown(metadata, reason):
	return _result(
		"UNKNOWN", "unclassified", None, "NONE", [reason], metadata,
		"Gather stronger disassembly/source evidence; do not infer a recipe or park class.")


def _ret_count(instructions):
	return sum(instruction["mnemonic"].startswith("ret") for instruction in instructions)


def _classify_return_shape(target, candidate, metadata):
	target_count = _ret_count(target)
	candidate_count = _ret_count(candidate)
	if target_count == candidate_count or not metadata["relocation_identities_equal"]:
		return None
	return _result(
		"RECIPE", "control-flow-return-shape", "R5", "PROVEN", [
			f"target ret count={target_count}; candidate ret count={candidate_count}",
			"ordered relocation identities agree",
		], metadata,
		"Stop after the return count matches; any residual needs a fresh classification.", [
			"Preserve the target's early-return versus single-exit structure.",
		])


def _has_frame(instructions):
	keys = [_instruction_key(instruction) for instruction in instructions[:8]]
	push_ebp = any(mnemonic.startswith("push") and operands == "%ebp"
					for mnemonic, operands in keys)
	mov_frame = any(mnemonic.startswith("mov") and operands == "%esp, %ebp"
					for mnemonic, operands in keys)
	sub_stack = any(mnemonic.startswith("sub") and operands.endswith(", %esp")
					for mnemonic, operands in keys)
	return push_ebp and mov_frame and sub_stack


def _classify_private_abi(target, candidate, metadata):
	if _has_frame(target) or not _has_frame(candidate):
		return None
	target_head = " ".join(instruction["operands"] for instruction in target[:12])
	candidate_head = " ".join(instruction["operands"] for instruction in candidate[:12])
	if not ("%ecx" in target_head or "%edx" in target_head):
		return None
	if not re.search(r"0x(?:8|c)\(%ebp\)", candidate_head):
		return None
	if not metadata["relocation_identities_equal"]:
		return None
	return _result(
		"RECIPE", "private-abi-frame-convention", "R1", "STRONG", [
			"target is frameless and consumes ECX/EDX near entry",
			"candidate builds an EBP frame and reads stack arguments",
			"ordered relocation identities agree",
		], metadata,
		"Stop if the helper still requires a stack frame after shortening live ranges.", [
			"Confirm the source helper is static.",
			"Confirm its address is never taken and it has a same-TU caller.",
		])


def _has_sib_scale_eight(instructions):
	for instruction in instructions:
		if not instruction["mnemonic"].startswith("mov"):
			continue
		operands = _split_operands(instruction["operands"])
		if operands and re.search(r"\([^)]*,[^)]*,8\)", operands[-1]):
			return True
	return False


def _has_scaled_temporary(instructions):
	return any(
		(instruction["mnemonic"].startswith("lea")
		 and re.search(r"\([^)]*,[^)]*,8\)", instruction["operands"]))
		or (instruction["mnemonic"].startswith("shl")
			and re.search(r"\$(?:0x)?3,", instruction["operands"]))
		for instruction in instructions)


def _classify_alias_shape(target, candidate, metadata):
	if not _has_sib_scale_eight(target) or not _has_scaled_temporary(candidate):
		return None
	return _result(
		"RECIPE", "alias-global-ownership", "R3", "PROVEN", [
			"target contains direct scale-8 SIB stores",
			"candidate materializes a scale-8 address with LEA/SHL",
		], metadata,
		"Stop and classify as ownership-blocked if target relocations use one aggregate anchor.", [
			"Audit target COFF ownership before splitting an aggregate global.",
			"After any source change, require strict relocation semantics as well as code bytes.",
		])


def _fst_kind(mnemonic):
	if mnemonic.startswith("fstp"):
		return "fstp", mnemonic[4:]
	if mnemonic.startswith("fst"):
		return "fst", mnemonic[3:]
	return None


def _classify_x87_liveness(target, candidate, metadata):
	if len(target) != len(candidate) or not metadata["relocation_identities_equal"]:
		return None
	differences = []
	for index, (target_instruction, candidate_instruction) in enumerate(zip(target, candidate)):
		if _instruction_key(target_instruction) == _instruction_key(candidate_instruction):
			continue
		target_kind = _fst_kind(target_instruction["mnemonic"])
		candidate_kind = _fst_kind(candidate_instruction["mnemonic"])
		if (not target_kind or not candidate_kind
				or target_kind[0] == candidate_kind[0]
				or target_kind[1] != candidate_kind[1]
				or target_instruction["operands"] != candidate_instruction["operands"]):
			return None
		differences.append((index, target_kind[0], candidate_kind[0]))
	if not differences:
		return None
	return _result(
		"RECIPE", "x87-value-liveness", "R4", "PROVEN", [
			"all instruction differences are FST/FSTP opcode pairs",
			"ordered relocation identities agree",
			*[f"instruction {index}: target {left}, candidate {right}"
			  for index, left, right in differences],
		], metadata,
		"Stop after x87 stack depth and subsequent instruction addresses realign.", [
			"Target FST requires keeping the computed value live after the store.",
			"Target FSTP requires letting the computed value die at the store.",
		])


def _replace_registers(operands, mapping):
	def replace(match):
		alias = match.group(1)
		register = ALIAS_TO_REGISTER[alias]
		mapped = mapping.get(register, register)
		aliases = REGISTER_ALIASES[mapped]
		original_aliases = REGISTER_ALIASES[register]
		try:
			index = original_aliases.index(alias)
		except ValueError:
			return match.group(0)
		if index >= len(aliases):
			return match.group(0)
		return "%" + aliases[index]
	return REGISTER_RE.sub(replace, operands)


def _classify_register_permutation(target, candidate, metadata):
	if (not metadata["size_equal"] or not metadata["relocations_exact"]
			or len(target) != len(candidate)):
		return None
	if [item["mnemonic"] for item in target] != [item["mnemonic"] for item in candidate]:
		return None
	target_keys = [_instruction_key(instruction) for instruction in target]
	candidate_keys = [_instruction_key(instruction) for instruction in candidate]
	differing = [
		index for index, pair in enumerate(zip(target_keys, candidate_keys))
		if pair[0] != pair[1]
	]
	if not differing:
		return None
	# A register live range can begin and end inside a function.  Permit one
	# contiguous local window, but refuse disconnected differences because a
	# single causal permutation is then unproven.
	if differing != list(range(differing[0], differing[-1] + 1)):
		return None
	window = range(differing[0], differing[-1] + 1)
	for permutation in itertools.permutations(GENERAL_REGISTERS):
		mapping = dict(zip(GENERAL_REGISTERS, permutation))
		if all(mapping[register] == register for register in GENERAL_REGISTERS):
			continue
		if all(
			_replace_registers(candidate_instruction["operands"], mapping)
			== target_instruction["operands"]
			for target_instruction, candidate_instruction in
			((target[index], candidate[index]) for index in window)
		):
			changed = [
				f"{source}->{destination}"
				for source, destination in mapping.items()
				if source != destination
			]
			return _result(
				"STOP", "pure-register-permutation", None, "PROVEN", [
					"equal padded size and exact relocation semantics",
					f"one contiguous instruction window {differing[0]}..{differing[-1]} "
					"becomes identical under one bijective GPR mapping",
					"mapping: " + ", ".join(changed),
				], metadata,
				"No measured legal-C control moved a pure register permutation; keep parked.")
	return None


def _registers(instruction):
	return {
		ALIAS_TO_REGISTER[match.group(1)]
		for match in REGISTER_RE.finditer(instruction["operands"])
	}


def _memory_role(instruction):
	operands = instruction["operands"]
	if "(" not in operands:
		return "none"
	mnemonic = instruction["mnemonic"]
	if mnemonic.startswith("lea"):
		return "none"
	parts = _split_operands(operands)
	if mnemonic.startswith("mov") and parts and "(" in parts[-1]:
		return "write"
	if mnemonic.startswith("fst"):
		return "write"
	return "read"


def _flag_role(mnemonic):
	if mnemonic.startswith(("j", "set", "cmov", "adc", "sbb")):
		return "read"
	if mnemonic.startswith((
		"add", "sub", "cmp", "test", "and", "or", "xor", "inc", "dec",
		"shl", "shr", "sar", "sal", "neg", "imul")):
		return "write"
	return "none"


def _independent(left, right):
	if left["mnemonic"].startswith(("j", "call", "ret", "loop", "int", "push", "pop")):
		return False
	if right["mnemonic"].startswith(("j", "call", "ret", "loop", "int", "push", "pop")):
		return False
	if left["mnemonic"].startswith("f") and right["mnemonic"].startswith("f"):
		return False
	if _registers(left) & _registers(right):
		return False
	left_memory = _memory_role(left)
	right_memory = _memory_role(right)
	if "write" in (left_memory, right_memory) and "none" not in (left_memory, right_memory):
		return False
	left_flags = _flag_role(left["mnemonic"])
	right_flags = _flag_role(right["mnemonic"])
	if left_flags == "read" or right_flags == "read":
		return False
	if left_flags == right_flags == "write":
		return False
	return True


def _classify_scheduling_tie(target, candidate, metadata):
	if (not metadata["size_equal"] or not metadata["relocations_exact"]
			or len(target) != len(candidate)):
		return None
	target_keys = [_instruction_key(instruction) for instruction in target]
	candidate_keys = [_instruction_key(instruction) for instruction in candidate]
	differing = [
		index for index, pair in enumerate(zip(target_keys, candidate_keys))
		if pair[0] != pair[1]
	]
	if len(differing) != 2 or differing[1] != differing[0] + 1:
		return None
	first, second = differing
	if not (target_keys[first] == candidate_keys[second]
			and target_keys[second] == candidate_keys[first]):
		return None
	if not _independent(target[first], target[second]):
		return None
	return _result(
		"STOP", "independent-instruction-scheduling", None, "PROVEN", [
			"equal padded size and exact relocation semantics",
			f"instructions {first} and {second} are one adjacent transposition",
			"the pair has disjoint registers, no control transfer, and no memory/flag dependency",
		], metadata,
		"No measured legal-C control moved an independent scheduling tie; keep parked.")


def _run_objdump(objdump, object_path, function_name):
	command = [
		str(objdump), "-d", "-r", "--symbolize-operands",
		f"--disassemble-symbols={function_name}", str(object_path),
	]
	completed = subprocess.run(
		command, capture_output=True, text=True, check=False)
	if completed.returncode:
		raise ClassifierError(completed.stderr.strip() or "llvm-objdump failed")
	instructions = parse_disassembly(completed.stdout)
	if not instructions:
		raise ClassifierError(f"no instructions found for {function_name!r}")
	return instructions


def evidence_from_object(object_path, function_name, objdump=None):
	"""Measure one COFF function and attach llvm-objdump disassembly evidence."""
	object_path = Path(object_path)
	try:
		measurement = section_info(load(object_path), function_name)
	except (CoffError, OSError) as error:
		raise ClassifierError(str(error)) from error
	objdump = objdump or shutil.which("llvm-objdump") or shutil.which("llvm-objdump.exe")
	if not objdump:
		raise ClassifierError("llvm-objdump was not found; pass --objdump")
	measurement["instructions"] = _run_objdump(objdump, object_path, function_name)
	return measurement


def main(argv=None):
	parser = argparse.ArgumentParser(description=__doc__)
	parser.add_argument("target")
	parser.add_argument("candidate")
	parser.add_argument("function")
	parser.add_argument("--objdump")
	parser.add_argument("--output", type=Path)
	args = parser.parse_args(argv)
	try:
		target = evidence_from_object(
			args.target, args.function, objdump=args.objdump)
		candidate = evidence_from_object(
			args.candidate, args.function, objdump=args.objdump)
		result = classify(target, candidate)
	except ClassifierError as error:
		parser.error(str(error))
	serialized = json.dumps(result, indent=2) + "\n"
	if args.output:
		args.output.parent.mkdir(parents=True, exist_ok=True)
		args.output.write_text(serialized, encoding="utf-8")
	print(serialized, end="")
	return 0


if __name__ == "__main__":
	raise SystemExit(main())