import wasmtime
import struct
import math
import os
import ctypes
from pathlib import Path


class Solver:
    def __init__(self):
        """Initialize the WASM solver"""
        # Try to find the WASM file
        current_dir = Path(__file__).parent
        go_wasm_path = Path(__file__).parent.parent.parent / 'deepseek4free' / 'pkg' / 'solver' / 'sha3_wasm_bg.7b9ca65ddd.wasm'

        if go_wasm_path.exists():
            with open(go_wasm_path, 'rb') as f:
                wasm_bytes = f.read()
        else:
            raise FileNotFoundError(f"WASM file not found at {go_wasm_path}")

        engine = wasmtime.Engine()
        self.module = wasmtime.Module(engine, wasm_bytes)
        self.store = wasmtime.Store(engine)
        self.linker = wasmtime.Linker(engine)
        self.linker.define_wasi()

        self.instance = self.linker.instantiate(self.store, self.module)
        
        # Get exports - handle both old and new wasmtime-py API
        exports = self.instance.exports(self.store)
        
        # Get memory export
        try:
            # Try direct attribute access first
            self.memory = exports.memory
        except AttributeError:
            # Try dict-like access
            try:
                self.memory = exports['memory']
            except (KeyError, TypeError):
                # Try get_export method
                mem_extern = self.instance.get_export(self.store, 'memory')
                if mem_extern and hasattr(mem_extern, 'memory'):
                    self.memory = mem_extern.memory
                else:
                    raise RuntimeError("Could not find memory export in WASM module")
        
        # Initialize functions - with error handling
        try:
            self.alloc_fn = exports.__wbindgen_export_0
        except AttributeError:
            self.alloc_fn = exports['__wbindgen_export_0']
        
        try:
            self.stack_ptr_fn = exports.__wbindgen_add_to_stack_pointer
        except AttributeError:
            self.stack_ptr_fn = exports['__wbindgen_add_to_stack_pointer']
        
        try:
            self.solve_fn = exports.wasm_solve
        except AttributeError:
            self.solve_fn = exports['wasm_solve']

    def _write_to_memory(self, text: str) -> tuple[int, int]:
        """Write a string to WASM memory and return pointer and length"""
        text_bytes = text.encode('utf-8')
        length = len(text_bytes)

        # Allocate memory - pass store as first argument
        ptr = self.alloc_fn(self.store, length, 1)
        print(f"[DEBUG] Allocated memory at ptr={ptr}, length={length}")

        # Get memory data pointer
        mem_ptr = self.memory.data_ptr(self.store)
        print(f"[DEBUG] Memory pointer type: {type(mem_ptr)}")
        
        # Write to memory - mem_ptr is a ctypes pointer, can index it directly
        try:
            for i, byte in enumerate(text_bytes):
                mem_ptr[ptr + i] = byte
            print(f"[DEBUG] Successfully wrote {length} bytes to memory")
        except TypeError as e:
            print(f"[DEBUG] Error writing to memory: {e}")
            # Try alternative approach with ctypes address
            try:
                addr = ctypes.cast(mem_ptr, ctypes.c_void_p).value
                print(f"[DEBUG] Memory address: {addr}")
                buffer = (ctypes.c_ubyte * length).from_address(addr + ptr)
                for i, byte in enumerate(text_bytes):
                    buffer[i] = byte
                print(f"[DEBUG] Successfully wrote {length} bytes using ctypes buffer")
            except Exception as e2:
                print(f"[DEBUG] Also failed with ctypes: {e2}")
                raise

        return ptr, length

    def _read_memory(self, ptr: int, length: int) -> bytes:
        """Read bytes from WASM memory"""
        mem_ptr = self.memory.data_ptr(self.store)
        print(f"[DEBUG] Reading {length} bytes from ptr={ptr}, memory_ptr type={type(mem_ptr)}")
        
        try:
            # Try direct indexing of ctypes pointer
            result = []
            for i in range(length):
                result.append(mem_ptr[ptr + i])
            return bytes(result)
        except TypeError as e:
            print(f"[DEBUG] Error reading with direct indexing: {e}")
            # Try converting to address and using ctypes
            try:
                addr = ctypes.cast(mem_ptr, ctypes.c_void_p).value
                print(f"[DEBUG] Memory address: {addr}")
                buffer = (ctypes.c_ubyte * length).from_address(addr + ptr)
                return bytes(buffer)
            except Exception as e2:
                print(f"[DEBUG] Also failed with ctypes: {e2}")
                raise

    def calculate_hash(self, challenge: str, salt: str, difficulty: int, expire_at: int) -> int:
        """Calculate hash using WASM solver"""
        print(f"[DEBUG] calculate_hash called with challenge={challenge[:20]}..., salt={salt}, difficulty={difficulty}")
        prefix = f"{salt}_{expire_at}_"

        # Adjust stack pointer - pass store as first argument
        retptr = self.stack_ptr_fn(self.store, -16)
        print(f"[DEBUG] Stack pointer adjusted, retptr={retptr}")

        # Write to memory
        challenge_ptr, challenge_len = self._write_to_memory(challenge)
        prefix_ptr, prefix_len = self._write_to_memory(prefix)
        print(f"[DEBUG] Wrote challenge at {challenge_ptr}, prefix at {prefix_ptr}")

        # Call solve function - pass store as first argument
        print(f"[DEBUG] Calling solve function with retptr={retptr}, challenge_ptr={challenge_ptr}, difficulty={difficulty}")
        self.solve_fn(self.store, retptr, challenge_ptr, challenge_len, prefix_ptr, prefix_len, float(difficulty))

        # Read result from memory
        status_bytes = self._read_memory(retptr, 4)
        status = struct.unpack('<I', status_bytes)[0]
        print(f"[DEBUG] Status: {status}")

        if status == 0:
            raise Exception("No solution found")

        # Read the answer as float64
        value_bytes = self._read_memory(retptr + 8, 8)
        value = struct.unpack('<d', value_bytes)[0]
        print(f"[DEBUG] Value from memory: {value}")

        # Convert float to int
        answer = int(value)
        print(f"[DEBUG] Final answer: {answer}")

        # Reset stack pointer - pass store as first argument
        self.stack_ptr_fn(self.store, 16)

        return answer

    def close(self):
        """Close the solver"""
        self.store = None
        self.instance = None