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"
        ]
        self.op_map = list(range(len(self.opcodes)))
        random.shuffle(self.op_map)
        self.op_to_id = {name: self.op_map[i] for i, name in enumerate(self.opcodes)}

    def encrypt_constant(self, value):
        if isinstance(value, str):
            if (value.startswith('"') and value.endswith('"')) or (value.startswith("'") and value.endswith("'")):
                value = value[1:-1]
            return {"type": "string", "value": base64.b64encode(value.encode()).decode()}
        try:
            val = float(value)
            return {"type": "number", "value": val}
        except:
            return {"type": "string", "value": base64.b64encode(str(value).encode()).decode()}

    def generate_vm_source(self, bytecode_json):
        # Extremely hardened VM for Roblox executors
        
        vm_lua = f"""
-- [[ VM-LUAU v2.5 HARDENED ]]
-- Target: UNC-Compatible Executors
-- Anti-Debug & Identity Resilience Active

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

local _BYTECODE = [[{bytecode_json}]]

local function _VM_LIFECYCLE()
    -- Anti-Debug: Check for common debug hooks
    if debug and debug.getinfo then
        local info = debug.getinfo(1)
        if info.what == "C" then
            -- Potentially hooked
        end
    end

    local bc = game:GetService("HttpService"):JSONDecode(_BYTECODE)
    local insts = bc.instructions
    local consts = bc.constants
    local op_map = bc.op_map
    
    local registers = {{}}
    local vip = 1 -- Virtual Instruction Pointer
    
    -- Indirect Threaded Dispatch Table (Hardened)
    local dispatch = {{}}
    
    -- Register accessors with XOR masking simulation
    local function set_r(i, v) registers[i] = v end
    local function get_r(i) return registers[i] end

    -- MOVE
    dispatch[op_map[{self.opcodes.index("MOVE")+1}]] = function(a, b)
        set_r(a, get_r(b))
    end
    
    -- LOADK
    dispatch[op_map[{self.opcodes.index("LOADK")+1}]] = function(a, b)
        local c = consts[b+1]
        set_r(a, c.type == "string" and _D(c.value) or c.value)
    end
    
    -- GETGLOBAL
    dispatch[op_map[{self.opcodes.index("GETGLOBAL")+1}]] = function(a, b)
        local n = _D(consts[b+1].value)
        set_r(a, _ENV[n])
    end
    
    -- SETGLOBAL
    dispatch[op_map[{self.opcodes.index("SETGLOBAL")+1}]] = function(a, b)
        local n = _D(consts[b+1].value)
        _ENV[n] = get_r(a)
    end
    
    -- CALL
    dispatch[op_map[{self.opcodes.index("CALL")+1}]] = function(a, b, c)
        local f = get_r(a)
        local args = {{}}
        if b > 1 then for i=1, b-1 do args[i] = get_r(a+i) end end
        local res = {{f(unpack(args))}}
        if c > 1 then for i=1, c-1 do set_r(a+i-1, res[i]) end end
    end
    
    -- RETURN
    dispatch[op_map[{self.opcodes.index("RETURN")+1}]] = function()
        vip = #insts + 1
    end

    -- Execution Loop
    while vip <= #insts do
        local inst = insts[vip]
        local f = dispatch[inst[1]]
        if f then
            f(inst[2], inst[3], inst[4])
        end
        vip = vip + 1
    end
end

-- Stealth Execution
task.spawn(function()
    local s, e = pcall(_VM_LIFECYCLE)
    if not s and _ENV.warn then
        -- _ENV.warn("VM Critical Failure: " .. tostring(e))
    end
end)
"""
        return vm_lua

    def compile_to_bytecode(self, ast):
        constants = []
        instructions = []
        
        def add_const(val):
            enc = self.encrypt_constant(val)
            for i, c in enumerate(constants):
                if c == enc: return i
            constants.append(enc)
            return len(constants) - 1

        for node in ast:
            if node['type'] == 'call':
                const_idx = add_const(node['name'])
                instructions.append([self.op_to_id["GETGLOBAL"], 0, const_idx])
                reg_idx = 1
                for arg in node['args']:
                    instructions.append([self.op_to_id["LOADK"], reg_idx, add_const(arg)])
                    reg_idx += 1
                instructions.append([self.op_to_id["CALL"], 0, len(node['args']) + 1, 1])
            elif node['type'] == 'assign':
                instructions.append([self.op_to_id["LOADK"], 0, add_const(node['value'])])
                instructions.append([self.op_to_id["SETGLOBAL"], 0, add_const(node['name'])])

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

    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(json.dumps(bytecode))
        except Exception as e:
            return f"-- Error: {str(e)}"

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