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):
            # Remove quotes if present
            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):
        # Hardened VM for Luau
        vm_lua = f"""
-- [[ VM-LUAU HARDENED ]]
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()
    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
    local dispatch = {{}}
    
    local function set_r(i, v) registers[i] = v end
    local function get_r(i) return registers[i] end

    dispatch[op_map[{self.opcodes.index("MOVE")+1}]] = function(a, b)
        set_r(a, get_r(b))
    end
    
    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
    
    dispatch[op_map[{self.opcodes.index("GETGLOBAL")+1}]] = function(a, b)
        local n = _D(consts[b+1].value)
        set_r(a, _ENV[n])
    end
    
    dispatch[op_map[{self.opcodes.index("SETGLOBAL")+1}]] = function(a, b)
        local n = _D(consts[b+1].value)
        _ENV[n] = get_r(a)
    end
    
    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
    
    dispatch[op_map[{self.opcodes.index("RETURN")+1}]] = function()
        vip = #insts + 1
    end

    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

task.spawn(_VM_LIFECYCLE)
"""
        return vm_lua

    def compile_to_bytecode(self, ast):
        constants = []
        instructions = []
        locals_map = {}
        next_reg = 0
        
        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

        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']]])
                else:
                    instructions.append([self.op_to_id["GETGLOBAL"], reg, add_const(expr['value'])])
            elif expr['type'] in ['STRING', 'NUMBER']:
                instructions.append([self.op_to_id["LOADK"], reg, add_const(expr['value'])])

        for node in ast:
            if node['type'] == 'call':
                func_reg = next_reg
                # func_reg might be used, so we need to be careful if we have many locals
                # For this simple obfuscator, we use registers above the current locals
                
                # Load function
                if node['name'] in locals_map:
                    instructions.append([self.op_to_id["MOVE"], func_reg, locals_map[node['name']]])
                else:
                    instructions.append([self.op_to_id["GETGLOBAL"], func_reg, add_const(node['name'])])
                
                # Load args
                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'])])

        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:
            import traceback
            return f"-- Error: {str(e)}\n{traceback.format_exc()}"

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