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

class LuauVMObfuscator:
    def __init__(self):
        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.k1 = random.randint(1000, 5000)
        self.k2 = random.randint(1000, 5000)
        
        self.op_to_id = {name: (self.opcodes.index(name) + self.k1) ^ self.k2 for name in self.opcodes}
        
        self.var_map = {}
        self.used_vars = set()

    def get_var(self, hint="var"):
        if hint in self.var_map: return self.var_map[hint]
        chars = "lI1" 
        length = random.randint(32, 48)
        new_var = "_" + "".join(random.choice(chars) for _ in range(length))
        while new_var in self.used_vars:
            new_var = "_" + "".join(random.choice(chars) for _ in range(length))
        self.used_vars.add(new_var)
        self.var_map[hint] = new_var
        return new_var

    def to_expr(self, n, bit32_var="bit32", depth=0):
        if depth > 2 or (depth > 0 and random.random() < 0.2):
            return str(n)
        
        r = random.randint(1, 100000)
        choice = random.choice(['add', 'sub', 'xor'])
        
        if choice == 'add':
            return f"({self.to_expr(n - r, bit32_var, depth + 1)} + {self.to_expr(r, bit32_var, depth + 1)})"
        elif choice == 'sub':
            return f"({self.to_expr(n + r, bit32_var, depth + 1)} - {self.to_expr(r, bit32_var, depth + 1)})"
        elif choice == 'xor':
            return f"{bit32_var}.bxor({self.to_expr(n ^ r, bit32_var, depth + 1)}, {self.to_expr(r, bit32_var, depth + 1)})"
        
        return str(n)

    def encrypt_string(self, s, key):
        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 minify(self, code):
        code = re.sub(r'--.*', '', code)
        code = re.sub(r'\s+', ' ', code)
        for op in ['==', '~=', '<=', '>=', '=', r'\+', '-', r'\*', '/', '>', '<', r'\(', r'\)', r'\{', r'\}', r'\[', r'\]', ',', ';', ':']:
            code = re.sub(r'\s*' + op + r'\s*', op.replace('\\', ''), code)
        return code.strip()

    def generate_vm_source(self, bytecode):
        raw_instructions = bytecode['instructions']
        indices = list(range(len(raw_instructions)))
        random.shuffle(indices)
        
        pos_map = {orig: shuffled_idx for shuffled_idx, orig in enumerate(indices)}
        
        final_insts = [None] * len(raw_instructions)
        for i, orig_idx in enumerate(indices):
            inst = raw_instructions[orig_idx]
            next_shuffled = pos_map[orig_idx + 1] if orig_idx + 1 < len(raw_instructions) else 0
            
            packed = [
                inst[0] & 0xFF, (inst[0] >> 8) & 0xFF, 
                inst[1] & 0xFF, inst[2] & 0xFF, inst[3] & 0xFF,
                next_shuffled & 0xFF, (next_shuffled >> 8) & 0xFF
            ]
            final_insts[i] = packed

        inst_str = "".join(chr(i) for inst in final_insts for i in inst)
        inst_b64 = base64.b64encode(inst_str.encode('latin-1')).decode()
        
        encrypted_consts = []
        salt = random.randint(100000, 999999)
        for i, c in enumerate(bytecode['constants']):
            if c['type'] == 'string':
                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)
        start_idx = pos_map[0]

        V_ENV = self.get_var("env")
        V_BIT = self.get_var("bit")
        V_B64 = self.get_var("b64")
        V_DEC = self.get_var("dec")
        V_INST = self.get_var("inst")
        V_CONSTS = self.get_var("consts")
        V_SALT = self.get_var("salt")
        V_EXEC = self.get_var("exec")
        V_REGS = self.get_var("regs")
        V_CURR = self.get_var("curr")
        V_RUN = self.get_var("run")
        V_GETC = self.get_var("getc")
        V_PTR = self.get_var("ptr")
        V_OP = self.get_var("op")
        V_A = self.get_var("a")
        V_B = self.get_var("b")
        V_C = self.get_var("c")
        V_UNP = self.get_var("unp")
        V_SPW = self.get_var("spw")
        V_RAW = self.get_var("raw")
        V_KEY = self.get_var("key")
        V_RES = self.get_var("res")
        V_LST = self.get_var("lst")

        vm_lua = f"""local {V_BIT}=bit32
local {V_ENV}=setmetatable({{}},{{__index=getfenv()}})
local {V_UNP}=unpack or table.unpack
local {V_SPW}=task and task.spawn or spawn
local {V_B64}='ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
local {V_DEC}=function(data)
    data=data:gsub('[^'..{V_B64}..'=]','')
    local res,bits={{}},''
    for i=1,#data do
        local c=data:sub(i,i)
        if c=='=' then break end
        local f={V_B64}:find(c)-1
        for j=5,0,-1 do bits=bits..({V_BIT}.extract(f,j,1)) end
    end
    for i=1,#bits-7,8 do
        local b=0
        for j=1,8 do b=b+(bits:sub(i+j-1,i+j-1)=='1' and 2^(8-j) or 0) end
        res[#res+1]=string.char(b)
    end
    return table.concat(res)
end
local {V_INST}={V_DEC}('{inst_b64}')
local {V_CONSTS}=game:GetService("HttpService"):JSONDecode([=[{consts_json}]=])
local {V_SALT}={self.to_expr(salt, V_BIT)}
local function {V_EXEC}()
    local {V_REGS}={{}}
    local {V_CURR}={self.to_expr(start_idx, V_BIT)}
    local {V_RUN}=true
    local function {V_GETC}(idx)
        if not idx then return nil end
        local c={V_CONSTS}[idx+1]
        if not c then return nil end
        if c.t=={self.to_expr(1, V_BIT)} then
            local {V_RAW}={V_DEC}(c.v)
            local {V_KEY}=(idx*{self.to_expr(149, V_BIT)}+{V_SALT})%256
            local {V_RES}={{}}
            local {V_LST}={V_KEY}%256
            for i=1,#{V_RAW} do
                local k=({V_KEY}+i+{V_LST}-1)%256
                local b=string.byte({V_RAW},i)
                local char={V_BIT}.bxor(b,k)
                {V_RES}[i]=string.char(char)
                {V_LST}=char
            end
            return table.concat({V_RES})
        end
        return c.v
    end
    while {V_RUN} do
        local {V_PTR}={V_CURR}*7+1
        local op_l=string.byte({V_INST},{V_PTR}) or 0
        local op_h=string.byte({V_INST},{V_PTR}+1) or 0
        local {V_A}=string.byte({V_INST},{V_PTR}+2) or 0
        local {V_B}=string.byte({V_INST},{V_PTR}+3) or 0
        local {V_C}=string.byte({V_INST},{V_PTR}+4) or 0
        local n_l=string.byte({V_INST},{V_PTR}+5) or 0
        local n_h=string.byte({V_INST},{V_PTR}+6) or 0
        {V_CURR}=n_l+({V_BIT}.lshift(n_h,8))
        local op_raw=op_l+({V_BIT}.lshift(op_h,8))
        local {V_OP}={V_BIT}.bxor(op_raw,{self.to_expr(self.k2, V_BIT)})-{self.to_expr(self.k1, V_BIT)}
        if {V_OP}=={self.opcodes.index('MOVE')} then
            {V_REGS}[{V_A}]={V_REGS}[{V_B}]
        elseif {V_OP}=={self.opcodes.index('LOADK')} then
            {V_REGS}[{V_A}]={V_GETC}({V_B})
        elseif {V_OP}=={self.opcodes.index('GETGLOBAL')} then
            local n={V_GETC}({V_B})
            if n then {V_REGS}[{V_A}]={V_ENV}[n] end
        elseif {V_OP}=={self.opcodes.index('SETGLOBAL')} then
            local n={V_GETC}({V_B})
            if n then {V_ENV}[n]={V_REGS}[{V_A}] end
        elseif {V_OP}=={self.opcodes.index('CALL')} then
            local f={V_REGS}[{V_A}]
            if f then
                local args={{}}
                if {V_B}>1 then for i=1,{V_B}-1 do args[i]={V_REGS}[{V_A}+i] end end
                local res={{f({V_UNP}(args))}}
                if {V_C}>1 then for i=1,{V_C}-1 do {V_REGS}[{V_A}+i-1]=res[i] end end
            end
        elseif {V_OP}=={self.opcodes.index('RETURN')} then
            {V_RUN}=false
        end
    end
end
{V_SPW}({V_EXEC})"""
        return self.minify(vm_lua)

    def compile_to_bytecode(self, ast):
        constants = []
        instructions = []
        locals_map = {}
        next_reg = 0
        
        def add_const(val):
            if isinstance(val, str):
                s_q = chr(39); d_q = chr(34)
                if (val.startswith(s_q) and val.endswith(s_q)) or (val.startswith(d_q) and val.endswith(d_q)):
                    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 ""
        try:
            lexer = Lexer(code)
            tokens = lexer.tokenize()
            parser = Parser(tokens)
            ast = parser.parse()
            if not ast: return ""
            bytecode = self.compile_to_bytecode(ast)
            return self.generate_vm_source(bytecode)
        except Exception:
            import traceback
            return f"-- Obfuscation Error: {traceback.format_exc()}"

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