"""Local-only (127.0.0.1) JSON-over-TCP server wrapping a DebugAgent. Never bind 0.0.0.0
here -- the MCP server is the only LAN-reachable piece (see spec Security section)."""

from __future__ import annotations

import base64
import json
import socketserver
import threading

from tools.havoc_debug_agent import DebugAgent
from tools.havoc_debug_protocol import decode_message, WatchpointLimitError

_OPS = {}


def _op(name):
    def register(fn):
        _OPS[name] = fn
        return fn
    return register


@_op("launch")
def _op_launch(agent: DebugAgent, params: dict) -> dict:
    agent.launch(params["path"], params["cwd"])
    return {}


@_op("set_breakpoint")
def _op_set_breakpoint(agent: DebugAgent, params: dict) -> dict:
    agent.set_breakpoint(int(params["addr"], 16), params.get("condition"))
    return {}


@_op("remove_breakpoint")
def _op_remove_breakpoint(agent: DebugAgent, params: dict) -> dict:
    agent.remove_breakpoint(int(params["addr"], 16))
    return {}


@_op("list_breakpoints")
def _op_list_breakpoints(agent: DebugAgent, params: dict) -> dict:
    return {"breakpoints": [hex(addr) for addr in agent.list_breakpoints()]}


@_op("set_watchpoint")
def _op_set_watchpoint(agent: DebugAgent, params: dict) -> dict:
    slot = agent.set_watchpoint(int(params["addr"], 16), params["size"], params["mode"])
    return {"slot": slot}


@_op("remove_watchpoint")
def _op_remove_watchpoint(agent: DebugAgent, params: dict) -> dict:
    agent.remove_watchpoint(params["slot"])
    return {}


@_op("list_watchpoints")
def _op_list_watchpoints(agent: DebugAgent, params: dict) -> dict:
    return {"watchpoints": {slot: hex(addr) for slot, addr in agent.list_watchpoints().items()}}


@_op("continue_execution")
def _op_continue_execution(agent: DebugAgent, params: dict) -> dict:
    return agent.continue_execution(params["timeout"])


@_op("read_memory")
def _op_read_memory(agent: DebugAgent, params: dict) -> dict:
    data = agent.read_memory(int(params["addr"], 16), params["length"])
    return {"data": base64.b64encode(data).decode("ascii")}


@_op("write_memory")
def _op_write_memory(agent: DebugAgent, params: dict) -> dict:
    agent.write_memory(int(params["addr"], 16), base64.b64decode(params["data"]))
    return {}


@_op("get_registers")
def _op_get_registers(agent: DebugAgent, params: dict) -> dict:
    return {"registers": {name: hex(value) for name, value in agent.get_registers().items()}}


@_op("screenshot")
def _op_screenshot(agent: DebugAgent, params: dict) -> dict:
    png_bytes = agent.screenshot()
    return {"png_base64": base64.b64encode(png_bytes).decode("ascii")}


@_op("get_status")
def _op_get_status(agent: DebugAgent, params: dict) -> dict:
    return agent.get_status()


class _Handler(socketserver.StreamRequestHandler):
    def handle(self) -> None:
        agent: DebugAgent = self.server.agent  # type: ignore[attr-defined]
        lock: threading.Lock = self.server.lock  # type: ignore[attr-defined]
        for line in self.rfile:
            if not line.strip():
                continue
            request = decode_message(line)
            op_name = request["op"]
            params = request.get("params", {})
            try:
                with lock:
                    if op_name not in _OPS:
                        raise ValueError(f"unknown op: {op_name!r}")
                    result = _OPS[op_name](agent, params)
                response = {"ok": True, "result": result}
            except WatchpointLimitError as exc:
                response = {"ok": False, "error": str(exc)}
            except Exception as exc:  # noqa: BLE001
                response = {"ok": False, "error": f"{type(exc).__name__}: {exc}"}
            self.wfile.write((json.dumps(response) + "\n").encode("utf-8"))


class _Server(socketserver.ThreadingTCPServer):
    allow_reuse_address = True
    daemon_threads = True


def run_agent_server(agent: DebugAgent, port: int) -> None:
    """Blocks forever, serving the local-only debug agent API on 127.0.0.1:port."""
    server = _Server(("127.0.0.1", port), _Handler)
    server.agent = agent  # type: ignore[attr-defined]
    server.lock = threading.Lock()  # type: ignore[attr-defined]
    server.serve_forever()


if __name__ == "__main__":
    import sys
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 47474
    run_agent_server(DebugAgent(), port)