import random
import string
import base64
import json
from .parser import Lexer, Parser

class LuauVMObfuscator:
    def __init__(self):
        # Full Luau-style Opcode list
        self.opcodes = [
            "MOVE", "LOADK", "LOADBOOL", "LOADNIL", "GETGLOBAL", "SETGLOBAL",
            "GETTABLE", "SETTABLE", "NEWTABLE", "SELF", "ADD", "SUB", "MUL",
            "DIV", "MOD", "POW", "UNM", "NOT", "LEN", "CONCAT", "JMP", "EQ",
            "LT", "LE", "TEST", "TESTSET", "CALL", "TAILCALL", "RETURN",
            "FORLOOP", "FORPREP", "TFORLOOP", "SETLIST", "CLOSE", "CLOSURE", "VARARG"
        ]
        
        # Arithmetic keys for opcode decoding
        self.k1 = random.randint(50, 200)
        self.k2 = random.randint(50, 200)
        self.k3 = random.randint(1, 10) # Multiplier/Step
        
        # Opcode to encoded ID: (real_index + k1) ^ k2
        self.op_to_id = {name: ((self.opcodes.index(name) + self.k1) ^ self.k2) % 256 for name in self.opcodes}

    def encrypt_string(self, s, key):
        # More complex XOR scheme: each byte depends on previous byte and index
        res = []
        last = key % 256
        for i, c in enumerate(s):
            k = (key + i + last) % 256
            last = ord(c)
            res.append(chr(ord(c) ^ k))
        return "".join(res)

    def generate_vm_source(self, bytecode):
        raw_instructions = bytecode['instructions']
        # Each instruction is [OP, A, B, C]
        
        # SHUFFLE AND FLATTEN: 
        # We will add a 'next' index to each instruction to break linear execution.
        # Format: [OP, A, B, C, NEXT_IDX_L, NEXT_IDX_H] (6 bytes)
        shuffled = []
        indices = list(range(len(raw_instructions)))
        random.shuffle(indices)
        
        # Map original index to shuffled index
        pos_map = {orig: shuffled_idx for shuffled_idx, orig in enumerate(indices)}
        
        # Prepare 6-byte instructions
        final_insts = [None] * len(raw_instructions)
        for i, orig_idx in enumerate(indices):
            inst = raw_instructions[orig_idx]
            
            # Find next shuffled index
            if orig_idx + 1 < len(raw_instructions):
                next_orig = orig_idx + 1
                next_shuffled = pos_map[next_orig]
            else:
                next_shuffled = 0 # End
            
            # Pack: [OP, A, B, C, Next_L, Next_H]
            packed = [
                inst[0], inst[1], inst[2], inst[3],
                next_shuffled & 0xFF, (next_shuffled >> 8) & 0xFF
            ]
            final_insts[i] = packed

        # Pack instructions into a string
        inst_str = "".join(chr(i) for inst in final_insts for i in inst)
        inst_b64 = base64.b64encode(inst_str.encode('latin-1')).decode()
        
        # Prepare constants
        encrypted_consts = []
        salt = random.randint(1000, 9999)
        for i, c in enumerate(bytecode['constants']):
            if c['type'] == 'string':
                # Encrypt with complex key
                key = (i * 149 + salt) % 256
                enc_val = self.encrypt_string(c['value'], key)
                encrypted_consts.append({"t": 1, "v": base64.b64encode(enc_val.encode('latin-1')).decode()})
            else:
                encrypted_consts.append({"t": 2, "v": c['value']})
        
        consts_json = json.dumps(encrypted_consts)
        
        # Starting shuffled index
        start_idx = pos_map[0]

        vm_lua = f"""
-- [[ LUAU-VM HARDENED - V2 ]]
local _ENV = getfenv()
local _B64 = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
local _D = function(data)
    data = string.gsub(data, '[^'.._B64..'=]', '')
    return (data:gsub('.', function(x)
        if (x == '=') then return '' end
        local r,f='',(_B64:find(x)-1)
        for i=6,1,-1 do r=r..(f%2^i-f%2^(i-1)>0 and '1' or '0') end
        return r;
    end):gsub('%d%d%d%d%d%d%d%d', function(x)
        local r=0
        for i=1,8 do r=r+(x:sub(i,i)=='1' and 2^(8-i) or 0) end
        return string.char(r)
    end))
end

-- Anti-Analysis Layer
local function _CHECK()
    -- Check for common hooks or debuggers
    local suspicious = {{ "getgenv", "getrenv", "getreg", "debug" }}
    for _, s in ipairs(suspicious) do
        if _ENV[s] then 
            -- Fake failure or subtle corruption
            return true 
        end
    end
    return false
end

local _INST_RAW = _D('{inst_b64}')
local _CONSTS = game:GetService("HttpService"):JSONDecode('{consts_json}')
local _SALT = {salt}

local function _EXECUTE()
    if _CHECK() then
        -- Insert bogus delay or crash if hooked
        for i=1, 100000 do local x = math.sin(i) end
    end

    local registers = {{}}
    local current = {start_idx}
    local running = true
    
    local function get_const(idx)
        local c = _CONSTS[idx + 1]
        if not c then return nil end
        if c.t == 1 then
            local raw = _D(c.v)
            local key = (idx * 149 + _SALT) % 256
            local res = ""
            local last = key % 256
            for i=1, #raw do
                local k = (key + i + last - 1) % 256
                local b = string.byte(raw, i)
                local char_code = bit32.bxor(b, k)
                res = res .. string.char(char_code)
                last = char_code
            end
            return res
        end
        return c.v
    end

    -- MEGA-DISPATCH LOOP (Control-Flow Flattening)
    while running do
        local ptr = current * 6 + 1
        if ptr > #_INST_RAW then break end
        
        local op_raw = string.byte(_INST_RAW, ptr)
        local a = string.byte(_INST_RAW, ptr + 1)
        local b = string.byte(_INST_RAW, ptr + 2)
        local c = string.byte(_INST_RAW, ptr + 3)
        local next_l = string.byte(_INST_RAW, ptr + 4)
        local next_h = string.byte(_INST_RAW, ptr + 5)
        
        current = next_l + (next_h * 256)
        
        -- Arithmetic Opcode Decoding
        local op = bit32.bxor(op_raw, {self.k2}) - {self.k1}
        
        if op == {self.opcodes.index('MOVE')} then
            registers[a] = registers[b]
        elseif op == {self.opcodes.index('LOADK')} then
            registers[a] = get_const(b)
        elseif op == {self.opcodes.index('GETGLOBAL')} then
            registers[a] = _ENV[get_const(b)]
        elseif op == {self.opcodes.index('SETGLOBAL')} then
            _ENV[get_const(b)] = registers[a]
        elseif op == {self.opcodes.index('CALL')} then
            local f = registers[a]
            local args = {{}}
            if b > 1 then for i=1, b-1 do args[i] = registers[a+i] end end
            local res = {{f(unpack(args))}}
            if c > 1 then for i=1, c-1 do registers[a+i-1] = res[i] end end
        elseif op == {self.opcodes.index('RETURN')} then
            running = false
        end
        
        -- Bogus execution path to confuse static analysis
        if op == -999 then
            running = false
        end
    end
end

task.spawn(_EXECUTE)
"""
        return vm_lua

    def compile_to_bytecode(self, ast):
        constants = []
        instructions = []
        locals_map = {}
        next_reg = 0
        
        def add_const(val):
            if isinstance(val, str):
                if (val.startswith("'") and val.endswith("'")) or (val.startswith("\"") and val.endswith("\"")):
                    val = val[1:-1]
            
            for i, c in enumerate(constants):
                if c['value'] == val: return i
            
            t = 'string' if isinstance(val, str) else 'number'
            constants.append({'type': t, 'value': val})
            return len(constants) - 1

        def load_expr_to_reg(expr, reg):
            if expr['type'] == 'IDENT':
                if expr['value'] in locals_map:
                    instructions.append([self.op_to_id["MOVE"], reg, locals_map[expr['value']], 0])
                else:
                    instructions.append([self.op_to_id["GETGLOBAL"], reg, add_const(expr['value']), 0])
            elif expr['type'] in ['STRING', 'NUMBER']:
                val = expr['value']
                if expr['type'] == 'NUMBER':
                    try: val = float(val)
                    except: pass
                instructions.append([self.op_to_id["LOADK"], reg, add_const(val), 0])

        for node in ast:
            if node['type'] == 'call':
                func_reg = next_reg
                if node['name'] in locals_map:
                    instructions.append([self.op_to_id["MOVE"], func_reg, locals_map[node['name']], 0])
                else:
                    instructions.append([self.op_to_id["GETGLOBAL"], func_reg, add_const(node['name']), 0])
                
                for i, arg_expr in enumerate(node['args']):
                    load_expr_to_reg(arg_expr, func_reg + 1 + i)
                
                instructions.append([self.op_to_id["CALL"], func_reg, len(node['args']) + 1, 1])
                
            elif node['type'] == 'assign':
                val_reg = next_reg
                load_expr_to_reg(node['value'], val_reg)
                
                if node.get('local'):
                    locals_map[node['name']] = val_reg
                    next_reg += 1
                else:
                    instructions.append([self.op_to_id["SETGLOBAL"], val_reg, add_const(node['name']), 0])

        instructions.append([self.op_to_id["RETURN"], 0, 0, 0])
        return {"instructions": instructions, "constants": constants}

    def obfuscate(self, code):
        if not code.strip(): return "-- No input"
        try:
            lexer = Lexer(code)
            tokens = lexer.tokenize()
            parser = Parser(tokens)
            ast = parser.parse()
            if not ast: return "-- VM Parser: No valid structures found."
            bytecode = self.compile_to_bytecode(ast)
            return self.generate_vm_source(bytecode)
        except Exception as e:
            import traceback
            return f"-- Error: {str(e)}\n{traceback.format_exc()}"

def obfuscate(code):
    return LuauVMObfuscator().obfuscate(code)