import argparse
import mmap
import os
import re
import shutil
import subprocess
import tempfile
from torch.utils.cpp_extension import CUDA_HOME


def run_cuobjdump(file_path):
    command = [f'{CUDA_HOME}/bin/cuobjdump', '-sass', file_path]
    result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
    assert result.returncode == 0
    return result.stdout


def extract_instructions(sass, instruction="FMNMX3"):
    lines = sass.splitlines()
    collected = []
    current_segment = []
    arch_name, func_name = 'N/A', 'N/A'

    i = 0
    while i < len(lines):
        line = lines[i]
        if 'code for' in line:
            arch_name = line.lstrip().removeprefix('code for ').rstrip()
        elif 'Function :' in line:
            func_name = line.strip().removeprefix('Function :')
        elif instruction in line:
            # Add instruction and the following line as a pair
            if i + 1 < len(lines):
                current_segment.append((line, lines[i + 1]))
                i += 1  # Skip the next line since we just processed it
        else:
            # End of current segment - save if it has at least 2 instruction pairs
            if len(current_segment) >= 2:
                collected.append((f'{arch_name}::{func_name}', current_segment))
            current_segment = []
        i += 1
    # Handle final segment
    if len(current_segment) >= 2:
        collected.append((f'{arch_name}::{func_name}', current_segment))

    if int(os.getenv('FA_PRINT_REG_REUSE', 0)):
        print(f'Found {len(collected)} {instruction} segments')
    return collected


def extract_hex_from_line(line):
    match = re.search(r'/\*\s*(0x[0-9a-fA-F]+)\s*\*/', line)
    assert match
    return int(match.group(1), 16)


def validate(m, offset, le_bytes, num_lines):
    assert len(le_bytes) == num_lines
    assert m[offset:offset + 16] == le_bytes[0]
    for i in range(1, num_lines):
        if m[offset + i * 16:offset + i * 16 + 16] != le_bytes[i]:
            return False
    return True


def parse_registers(line):
    line = re.sub(r'/\*.*?\*/', '', line)
    line = line.replace(';', '')
    tokens = line.strip().split(',')
    registers = []
    for token in tokens:
        token = token.strip()
        words = token.split()
        for word in words:
            if word.startswith('R'):
                reg = word.split('.')[0]
                registers.append(reg)
    return registers


FLINE_RE = re.compile(r'\s*/\*\w{4}\*/\s*([^;]*;)\s*/\* 0x(\w{16}) \*/\s*')
SLINE_RE = re.compile(r'\s*/\* 0x(\w{16}) \*/\s*')
FNAME_RE = re.compile(r'\s*Function : (\w+)\s*')
BRA_RE = re.compile(r'(.*BRA(?:\.U)? )(0x\w+);')


def parseCtrl(sline):
    enc = int(SLINE_RE.match(sline).group(1), 16)
    stall = (enc >> 41) & 0xf
    yld = (enc >> 45) & 0x1
    wrtdb = (enc >> 46) & 0x7
    readb = (enc >> 49) & 0x7
    watdb = (enc >> 52) & 0x3f

    yld_str = 'Y' if yld == 0 else '-'
    wrtdb_str = '-' if wrtdb == 7 else str(wrtdb)
    readb_str = '-' if readb == 7 else str(readb)
    watdb_str = '--' if watdb == 0 else f'{watdb:02d}'
    return f'{watdb_str}:{readb_str}:{wrtdb_str}:{yld_str}:{stall:x}'


def modify_segment(m, name, instruction_lines):
    num_lines = len(instruction_lines)

    le_bytes, new_le_bytes = [], []
    reused_list = []
    num_changed = 0
    for i in range(num_lines):
        low_line, high_line = instruction_lines[i]
        dst_reg = parse_registers(low_line)[-2]
        low_hex, high_hex = extract_hex_from_line(low_line), extract_hex_from_line(high_line)
        le_bytes.append(low_hex.to_bytes(8, 'little') + high_hex.to_bytes(8, 'little'))
        # If not the last instruction, set yield flag to "no yield" (bit 45 = 1)
        if not (i == num_lines - 1):
            # Check if yield flag is currently set to yield (bit 45 = 0)
            currently_yields = (high_hex & (1 << 45)) == 0
            if currently_yields:
                # Set bit 45 to 1 (no yield)
                high_hex |= (1 << 45)
                # Do we need to set the reused flag for the last argument (for FFMA2)
                reused_cur = (high_hex & (1 << 60)) != 0
                dst_reg_cur = parse_registers(instruction_lines[i][0])[-1]
                dst_reg_next = parse_registers(instruction_lines[i + 1][0])[-1]
                if not reused_cur and dst_reg_cur == dst_reg_next:
                    # Set the reused flag (bit 60 = 1)
                    high_hex |= (1 << 60)
                    reused_list.append(i)
                num_changed += 1
        # if num_changed > 0 and i == num_lines - 1:
        #     currently_yields = (high_hex & (1 << 45)) == 0
        #     if not currently_yields:  # yield in the last instruction if we have changed
        #         high_hex &= ~(1 << 45)
        #         num_changed += 1
        new_le_bytes.append(low_hex.to_bytes(8, 'little') + high_hex.to_bytes(8, 'little'))
    if int(os.getenv('FA_PRINT_REG_REUSE', 0)):
        print(f' > segment `{name}` new reused list ({num_changed} changed): {reused_list}')

    # Find the offset
    offsets = []
    offset = m.find(le_bytes[0])
    while offset != -1:
        offsets.append(offset)
        offset = m.find(le_bytes[0], offset + 1)
    offsets = list(filter(lambda x: validate(m, x, le_bytes, num_lines), offsets))

    # Replace with `new_le_bytes`
    for offset in offsets:
        for i in range(num_lines):
            m[offset + i * 16:offset + i * 16 + 16] = new_le_bytes[i]


def process_cubin_bytes(cubin_data, instruction="FMNMX3"):
    """
    Process CUBIN bytes data, modify FMNMX3 yield flags, and return modified bytes.

    Args:
        cubin_data (bytes): Original CUBIN file data

    Returns:
        bytes: Modified CUBIN file data
    """
    assert instruction in ['FMNMX3', 'FFMA2', 'FADD'], "Only FMNMX3, FFMA2, FADD instructions are supported"
    with tempfile.NamedTemporaryFile(suffix='.cubin', delete=False) as tmp_file:
        tmp_path = tmp_file.name
        tmp_file.write(cubin_data)
        tmp_file.flush()
        try:
            # Get SASS from the temporary file
            output = run_cuobjdump(tmp_path)
            segments = extract_instructions(output, instruction)
            # Modify the temporary file
            with open(tmp_path, 'r+b') as f:
                mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE)
                for segment in segments:
                    modify_segment(mm, *segment)
                mm.close()
            # Read the modified data
            with open(tmp_path, 'rb') as f:
                modified_data = f.read()
            return modified_data
        finally:
            # Clean up temporary file
            os.unlink(tmp_path)


def process(path):
    if int(os.getenv('FA_PRINT_REG_REUSE', 0)):
        print(f'Processing {path}')
    # Create output path
    output_path = path.replace('.cubin', '_modified.cubin')
    # Copy original to output
    shutil.copy2(path, output_path)
    # Process the copy
    output = run_cuobjdump(output_path)
    segments = extract_instructions(output)
    with open(output_path, 'r+b') as f:
        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_WRITE)
        for segment in segments:
            modify_segment(mm, *segment)
        mm.close()
    print(f'Modified CUBIN saved to: {output_path}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Change yield flag of FMNMX3 instructions')
    parser.add_argument('--path', help='Path to the cubin file')
    args = parser.parse_args()

    process(args.path)
