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, self.k2, self.k3 = random.randint(500, 2000), random.randint(500, 2000), random.randint(500, 2000)
        self.op_to_id = {name: ((self.opcodes.index(name) + self.k1) ^ self.k2) + self.k3 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]
        c = "Il1"; l = random.randint(24, 32)
        v = "_" + "".join(random.choice(c) for _ in range(l))
        while v in self.used_vars: v = "_" + "".join(random.choice(c) for _ in range(l))
        self.used_vars.add(v); self.var_map[hint] = v
        return v

    def to_expr(self, n):
        r = random.randint(1, 1000)
        ch = random.choice(['add', 'sub', 'xor'])
        if ch == 'add': return f"({n-r}+{r})"
        if ch == 'sub': return f"({n+r}-{r})"
        return f"bit32.bxor({n^r},{r})"

    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_ins = bytecode['instructions']
        indices = list(range(len(raw_ins)))
        random.shuffle(indices); pos_map = {orig: i for i, orig in enumerate(indices)}
        final_insts = []
        for orig_idx in indices:
            inst = raw_ins[orig_idx]
            next_sh = pos_map[orig_idx + 1] if orig_idx + 1 < len(raw_ins) else 0
            final_insts.append([inst[0]&0xFF,(inst[0]>>8)&0xFF,inst[1]&0xFF,inst[2]&0xFF,inst[3]&0xFF,next_sh&0xFF,(next_sh>>8)&0xFF])
        inst_str = "".join(chr(i) for inst in final_insts for i in inst)
        inst_b64 = base64.b64encode(inst_str.encode('latin-1')).decode()
        
        consts = []
        salt = random.randint(100000, 999999)
        for i, c in enumerate(bytecode['constants']):
            if c['type'] == 'string':
                k = (i * 149 + salt) % 256; enc = self.encrypt_string(c['value'], k)
                consts.append({"t": 1, "v": base64.b64encode(enc.encode('latin-1')).decode()})
            elif c['type'] == 'number': consts.append({"t": 2, "v": c['value']})
            else: consts.append({"t": 3, "v": c['value']})
        
        v = {
            "BIT": "bit32", "ENV": self.get_var("env"), "DEC": self.get_var("dec"),
            "INST": self.get_var("inst"), "CONSTS": self.get_var("consts"), "SALT": self.to_expr(salt),
            "EXEC": self.get_var("exec"), "REGS": self.get_var("regs"), "CURR": self.get_var("curr"),
            "PTR": self.get_var("ptr"), "OP": self.get_var("op"), "A": self.get_var("a"),
            "B": self.get_var("b"), "C": self.get_var("c"), "UNP": "unpack or table.unpack",
            "SPW": "task and task.spawn or spawn", "JSON": "game:GetService('HttpService')",
            "CHARS": self.get_var("chars"), "LOOKUP": self.get_var("lookup"), "GETC": self.get_var("getc")
        }

        vm_lua = f"""local {v['ENV']}=setmetatable({{}},{{__index=getfenv()}})
local {v['CHARS']}='ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/'
local {v['LOOKUP']}={{}}
for i=1,64 do {v['LOOKUP']}[{v['CHARS']}:sub(i,i)]=i-1 end
local function {v['DEC']}(data)
    data=data:gsub('[^%a%d%+/]','')
    local res={{}}
    for i=1,#data,4 do
        local a,b,c,d={v['LOOKUP']}[data:sub(i,i)],{v['LOOKUP']}[data:sub(i+1,i+1)],{v['LOOKUP']}[data:sub(i+2,i+2)],{v['LOOKUP']}[data:sub(i+3,i+3)]
        local n={v['BIT']}.lshift(a or 0,18)+{v['BIT']}.lshift(b or 0,12)+{v['BIT']}.lshift(c or 0,6)+(d or 0)
        res[#res+1]=string.char({v['BIT']}.extract(n,16,8))
        if c then res[#res+1]=string.char({v['BIT']}.extract(n,8,8)) end
        if d then res[#res+1]=string.char({v['BIT']}.extract(n,0,8)) end
    end
    return table.concat(res)
end
local {v['INST']}={v['DEC']}('{inst_b64}')
local {v['CONSTS']}={v['JSON']}:JSONDecode([=[{json.dumps(consts)}]=])
local {v['EXEC']}=function()
    local {v['REGS']}={{}}
    local {v['CURR']}={self.to_expr(pos_map[0])}
    local function {v['GETC']}(idx)
        local c={v['CONSTS']}[idx+1]
        if not c then return end
        if c.t==1 then
            local r,k={v['DEC']}(c.v),(idx*149+{v['SALT']})%256
            local res,lst={{}},k%256
            for i=1,#r do
                local char={v['BIT']}.bxor(string.byte(r,i),(k+i+lst-1)%256)
                res[i]=string.char(char)
                lst=char
            end
            return table.concat(res)
        elseif c.t==3 then return c.v=='true' and true or (c.v=='false' and false or nil) end
        return c.v
    end
    while true do
        local {v['PTR']}={v['CURR']}*7+1
        local b1,b2,b3,b4,b5,b6,b7=string.byte({v['INST']},{v['PTR']},{v['PTR']}+6)
        if not b1 then break end
        {v['CURR']}=b6+({v['BIT']}.lshift(b7,8))
        local {v['OP']}={v['BIT']}.bxor(b1+{v['BIT']}.lshift(b2,8),{self.to_expr(self.k2)})-{self.to_expr(self.k1)}
        local {v['A']},{v['B']},{v['C']}=b3,b4,b5
        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 {v['REGS']}[{v['A']}]={v['ENV']}[{v['GETC']}({v['B']})]
        elseif {v['OP']}=={self.opcodes.index('SETGLOBAL')} then {v['ENV']}[{v['GETC']}({v['B']})]={v['REGS']}[{v['A']}]
        elseif {v['OP']}=={self.opcodes.index('GETTABLE')} then 
            local b={v['REGS']}[{v['B']}]
            if b then {v['REGS']}[{v['A']}]=b[{v['REGS']}[{v['C']}] or {v['GETC']}({v['C']})] end
        elseif {v['OP']}=={self.opcodes.index('SETTABLE')} then
            local a={v['REGS']}[{v['A']}]
            if a then a[{v['REGS']}[{v['B']}] or {v['GETC']}({v['B']})]={v['REGS']}[{v['C']}] end
        elseif {v['OP']}=={self.opcodes.index('ADD')} then {v['REGS']}[{v['A']}]=({v['REGS']}[{v['B']}] or 0)+({v['REGS']}[{v['C']}] or 0)
        elseif {v['OP']}=={self.opcodes.index('SUB')} then {v['REGS']}[{v['A']}]=({v['REGS']}[{v['B']}] or 0)-({v['REGS']}[{v['C']}] or 0)
        elseif {v['OP']}=={self.opcodes.index('MUL')} then {v['REGS']}[{v['A']}]=({v['REGS']}[{v['B']}] or 0)*({v['REGS']}[{v['C']}] or 0)
        elseif {v['OP']}=={self.opcodes.index('DIV')} then {v['REGS']}[{v['A']}]=({v['REGS']}[{v['B']}] or 0)/({v['REGS']}[{v['C']}] or 1)
        elseif {v['OP']}=={self.opcodes.index('UNM')} then {v['REGS']}[{v['A']}]=-({v['REGS']}[{v['B']}] or 0)
        elseif {v['OP']}=={self.opcodes.index('NOT')} then {v['REGS']}[{v['A']}]=not {v['REGS']}[{v['B']}]
        elseif {v['OP']}=={self.opcodes.index('LEN')} then {v['REGS']}[{v['A']}]=#{v['REGS']}[{v['B']}] or 0
        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 break end
    end
end
{v['SPW']}({v['EXEC']})"""
        return self.minify(vm_lua)

    def compile_to_bytecode(self, ast):
        constants, instructions, locals_map = [], [], {}
        self.next_reg = 0
        def add_const(val):
            if isinstance(val, str) and ((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'
            if val in ['true', 'false', 'nil']: t = 'keyword'
            constants.append({'type': t, 'value': val})
            return len(constants) - 1
        def emit(op, a=0, b=0, c=0): instructions.append([self.op_to_id[op], a, b, c])
        def gen_expr(expr, target):
            if not expr: return
            if expr['type'] == 'NUMBER': emit("LOADK", target, add_const(float(expr['value'])))
            elif expr['type'] == 'STRING': emit("LOADK", target, add_const(expr['value']))
            elif expr['type'] == 'KEYWORD': emit("LOADK", target, add_const(expr['value']))
            elif expr['type'] == 'variable':
                if expr['name'] in locals_map: emit("MOVE", target, locals_map[expr['name']])
                else: emit("GETGLOBAL", target, add_const(expr['name']))
            elif expr['type'] == 'index':
                br = self.next_reg; self.next_reg += 1; gen_expr(expr['base'], br)
                kr = self.next_reg; self.next_reg += 1; gen_expr(expr['key'], kr)
                emit("GETTABLE", target, br, kr); self.next_reg -= 2
            elif expr['type'] == 'binary':
                lr = self.next_reg; self.next_reg += 1; gen_expr(expr['left'], lr)
                rr = self.next_reg; self.next_reg += 1; gen_expr(expr['right'], rr)
                op_m = {'+': 'ADD', '-': 'SUB', '*': 'MUL', '/': 'DIV'}
                emit(op_m.get(expr['op'], 'ADD'), target, lr, rr); self.next_reg -= 2
            elif expr['type'] == 'unary':
                or_reg = self.next_reg; self.next_reg += 1; gen_expr(expr['operand'], or_reg)
                op_m = {'-': 'UNM', '#': 'LEN', 'not': 'NOT'}
                emit(op_m.get(expr['op'], 'UNM'), target, or_reg); self.next_reg -= 1
            elif expr['type'] == 'call': gen_call(expr, target)
        def gen_call(node, target):
            fr = self.next_reg; self.next_reg += 1
            if node['func']['type'] == 'variable':
                if node['func']['name'] in locals_map: emit("MOVE", fr, locals_map[node['func']['name']])
                else: emit("GETGLOBAL", fr, add_const(node['func']['name']))
            else: gen_expr(node['func'], fr)
            for i, arg in enumerate(node['args']):
                ar = self.next_reg; self.next_reg += 1; gen_expr(arg, ar)
            emit("CALL", fr, len(node['args']) + 1, 2)
            emit("MOVE", target, fr)
            self.next_reg = fr
        for node in ast:
            if node['type'] == 'assign':
                vr = self.next_reg; self.next_reg += 1; gen_expr(node['value'], vr)
                if node.get('local'): locals_map[node['name']] = vr
                elif 'target' in node:
                    t = node['target']
                    if t['type'] == 'index':
                        br = self.next_reg; self.next_reg += 1; gen_expr(t['base'], br)
                        kr = self.next_reg; self.next_reg += 1; gen_expr(t['key'], kr)
                        emit("SETTABLE", br, kr, vr); self.next_reg -= 2
                else: emit("SETGLOBAL", vr, add_const(node['name']))
            elif node['type'] == 'call': gen_call(node, self.next_reg)
        emit("RETURN")
        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)