import json
from enum import Enum

from metalware_sdk.havoc_client import HavocClient

class WatchType(Enum):
  READ = "read"
  WRITE = "write"

class ReplayDebugger:
  def __init__(self, client: HavocClient, run_id: str, testcase_id: str):
    """Initialize a replay debugger for a run.

    Args:
      client: HavocClient instance
      run_id: The run UUID (string)
      testcase_id: The testcase ID (e.g., "crash_0x0_unknown")
    """
    self._client = client
    self._run_id = run_id
    self._testcase_id = testcase_id

    self._client.start_debug_session(self._run_id, self._testcase_id)

  def _send_command(self, command: dict):
    result = self._client.send_debug_command(self._run_id, self._testcase_id, json.dumps(command))
    return json.loads(result)

  def run(self) -> str:
    res = self._send_command({"c": "continue"})
    if 'data' in res and res['data'] is not None and 'exit_reason' in res['data']: return res['data']['exit_reason']
    else: raise RuntimeError(f"Failed to run: {res}")

  def add_breakpoint(self, address: int):
    self._send_command({"c": "add_breakpoint", "address": address})

  def remove_breakpoint(self, address: int):
    self._send_command({"c": "remove_breakpoint", "address": address})

  def add_watchpoint(self, address: int, watch_type: WatchType) -> None:
    result = self._send_command({"c": "add_watchpoint", "address": address, "watch_type": watch_type.value})
    if 'success' in result and result['success']: return
    else: raise RuntimeError(f"Failed to add watchpoint: {result}")

  def remove_watchpoint(self, address: int, watch_type: WatchType):
    self._send_command({"c": "remove_watchpoint", "address": address, "watch_type": watch_type.value})

  def step(self) -> str:
    res = self._send_command({"c": "step_instruction"})
    if 'data' in res and res['data'] is not None and 'exit_reason' in res['data']: return res['data']['exit_reason']
    else: raise RuntimeError(res['message'])

  def step_back(self) -> str:
    res = self._send_command({"c": "reverse_step_instruction"})
    if 'data' in res and res['data'] is not None and 'exit_reason' in res['data']: return res['data']['exit_reason']
    else: raise RuntimeError(res['message'])

  def step_block(self) -> str:
    res = self._send_command({"c": "step_block"})
    if 'data' in res and res['data'] is not None and 'exit_reason' in res['data']: return res['data']['exit_reason']
    else: raise RuntimeError(res['message'])

  def step_back_block(self) -> str:
    res = self._send_command({"c": "reverse_step_block"})
    if 'data' in res and res['data'] is not None and 'exit_reason' in res['data']: return res['data']['exit_reason']
    else: raise RuntimeError(res['message'])

  def state(self) -> dict:
    result = self._send_command({"c": "state"})
    if 'data' in result: return result['data']
    else: raise RuntimeError(result['message'])

  def read_register(self, register_name: str) -> int:
    result = self._send_command({"c": "read_reg", "reg_name": register_name})
    if 'data' in result and 'value' in result['data']: return result['data']['value']
    else: raise RuntimeError(result['message'])

  def write_register(self, register_name: str, value: int) -> None:
    result = self._send_command({"c": "write_reg", "reg_name": register_name, "value": value})
    if 'success' in result and result['success']: return
    else: raise RuntimeError(result['message'])

  def list_registers(self) -> dict:
    result = self._send_command({"c": "list_regs"})
    if 'data' in result and 'registers' in result['data']: return result['data']['registers']
    else: raise RuntimeError(result['message'])

  def read_memory(self, address: int, size: int) -> bytes:
    if not isinstance(address, int): raise TypeError(f"Address must be int, got {type(address)}")
    if not isinstance(size, int): raise TypeError(f"Size must be int, got {type(size)}")
    if size > 0x1000: raise RuntimeError("Cannot read more than 4096 bytes at once")
    result = self._send_command({"c": "read_mem", "address": address, "size": size})
    if 'data' in result and result['data'] is not None: return bytes(result['data'])
    else: raise RuntimeError(result['message'])

  def write_memory(self, address: int, data: bytes) -> None:
    if not isinstance(data, bytes): raise TypeError(f"Data must be bytes, got {type(data)}")
    if len(data) > 0x1000: raise RuntimeError("Cannot write more than 4096 bytes at once")
    print("Data as hex: ", data.hex())
    result = self._send_command({"c": "write_mem", "address": address, "data": data.hex()})
    if 'success' in result and result['success']: return
    else: raise RuntimeError(result['message'])

  def disassemble(self, count: int = 10) -> list[str]:
    result = self._send_command({"c": "disassemble", "count": count})
    if 'data' in result and result['data'] is not None and 'disassembly' in result['data']: return result['data']['disassembly']
    else: raise RuntimeError(result['message'])

  def print_asm(self):
    asm = self.disassemble()
    current_pc = self.read_register("pc")
    print("asm: ", asm)
    for pc, disasm in asm:
      if pc == current_pc: print(f"> {hex(pc)}: {disasm}")
      else: print(f"{hex(pc)}: {disasm}")
    print("done")

  def list_watchpoints(self) -> list[tuple[int, WatchType]]:
    result = self._send_command({"c": "list_watchpoints"})
    if 'data' in result and 'watchpoints' in result['data']:
      return [(int(wp[0]), WatchType(wp[1])) for wp in result['data']['watchpoints']]
    else: raise RuntimeError(result['message'])

  def list_breakpoints(self) -> list[int]:
    result = self._send_command({"c": "list_breakpoints"})
    if 'data' in result and 'breakpoints' in result['data']: return result['data']['breakpoints']
    else: raise RuntimeError(result['message'])

  def backtrace(self) -> list[int]:
    result = self._send_command({"c": "backtrace"})
    if 'data' in result and 'backtrace' in result['data']: return result['data']['backtrace']
    else: raise RuntimeError(result['message'])

  def print_backtrace(self):
    backtrace = self.backtrace()
    for pc in backtrace:
      print(f"{hex(pc)}")

  def rewind(self) -> None:
    result = self._send_command({"c": "rewind"})
    if 'success' in result and result['success']: return
    else: raise RuntimeError(result['message'])

  def disassemble_range(self, start_addr: int, count: int) -> list[str]:
    result = self._send_command({"c": "disassemble_range", "start_addr": start_addr, "count": count})
    if 'data' in result and 'disassembly' in result['data']: return result['data']['disassembly']
    else: raise RuntimeError(result['message'])

  def get_input(self) -> list[dict]:
    """Get all MMIO streams with their current data and cursor positions.

    Returns a list of dicts with keys: address, cursor, length, data, sizes
    """
    result = self._send_command({"c": "get_input"})
    if 'data' in result and 'streams' in result['data']: return result['data']['streams']
    else: raise RuntimeError(result['message'])

  def write_input(self, streams: dict[int, bytes]) -> dict:
    """Write data for the entire input (all streams).

    Args:
      streams: A dict mapping stream addresses (int) to stream data (bytes)
               Example: {0x400e0468: b'\\x01\\x00\\x00\\x00', 0x400e0430: b'\\x38\\x3c\\x6e\\x71'}

    Returns a dict with keys: stream_count, total_bytes

    Raises RuntimeError if input is not forked.
    """
    # Convert to array of {addr, data} objects for JSON serialization
    streams_arr = [{"addr": addr, "data": data.hex()} for addr, data in streams.items()]
    result = self._send_command({"c": "write_input", "streams": streams_arr})
    if 'success' in result and result['success'] and 'data' in result and result['data'] is not None:
      return result['data']
    else: raise RuntimeError(result['message'])

  def fork_input(self, name: str) -> str:
    """Fork the current input to a new file in the forked directory.

    Args:
      name: The name for the forked input file. This will be placed in the
            `forked` subdirectory of the run directory.

    Returns:
      The testcase ID of the forked input (forked_{name}).

    Raises RuntimeError if fork fails.
    """
    result = self._send_command({"c": "fork_input", "name": name})
    if 'success' in result and result['success']:
      return f"forked_{name}"
    else: raise RuntimeError(result['message'])