"""Shared reverse-engineering lookups for the agents, backed by the project's own image and annotations.

Configured per project in a [reverse] table:
    image = 'R:/.../image.bin'        flat image of the executable, the first byte at `base`
    base = 0x401000
    facts = 'R:/.../facts.jsonl'      one JSON object per function: address, size, name, calls, data, ...
    imports = 'R:/.../imports.json'   optional: [{module, name, slot}]
    annotations = 'R:/.../annotations' names.csv, functions.csv, globals.csv, prototypes.h
    src = 'R:/.../src'                optional: the decompiled C files (<ADDR>_<Name>.c)

Every answer comes from those files, read fresh when they change, so the agents see the team's current names.
Nothing here writes anything."""
import csv
import json
import os
import re
import struct
from pathlib import Path

from .util import HarborError

MAX_BYTES = 0x4000


def _int(value):
    if isinstance(value, int):
        return value
    text = str(value).strip()
    return int(text, 16) if text.lower().startswith("0x") or re.fullmatch(r"[0-9A-Fa-f]{6,8}", text) else int(text)


class Reverse:
    def __init__(self, config):
        if not config or not config.get("image"):
            raise HarborError("this project has no [reverse] table (image, base, facts, annotations)")
        self.image = Path(config["image"])
        self.base = _int(config.get("base", 0x401000))
        self.facts_path = Path(config["facts"]) if config.get("facts") else None
        self.imports_path = Path(config["imports"]) if config.get("imports") else None
        self.annotations = Path(config["annotations"]) if config.get("annotations") else None
        self.src = Path(config["src"]) if config.get("src") else None
        self.conventions_path = Path(config["conventions"]) if config.get("conventions") else None
        self._cache = {}

    # ----------------------------------------------------------------- loading (cached by file mtime)
    def _load(self, key, path, loader):
        if path is None or not path.exists():
            return loader(None)
        stamp = path.stat().st_mtime
        hit = self._cache.get(key)
        if hit and hit[0] == stamp:
            return hit[1]
        value = loader(path)
        self._cache[key] = (stamp, value)
        return value

    def _bytes(self):
        return self._load("image", self.image, lambda p: p.read_bytes() if p else b"")

    def _facts(self):
        def load(p):
            out = {}
            if p:
                for line in p.read_text(encoding="utf-8").splitlines():
                    if line.strip():
                        d = json.loads(line)
                        out[_int(d["address"])] = d
            return out
        return self._load("facts", self.facts_path, load)

    def _callers(self):
        facts = self._facts()
        key = ("callers", id(facts))
        if key not in self._cache:
            callers = {}
            for addr, d in facts.items():
                for c in d.get("calls", []):
                    callers.setdefault(_int(c), []).append(addr)
            self._cache[key] = callers
        return self._cache[key]

    def _imports(self):
        return self._load("imports", self.imports_path, lambda p: {
            int(d["slot"]): f"{d['module']}!{d['name']}" for d in json.loads(p.read_text(encoding="utf-8"))} if p else {})

    def _csv(self, name):
        path = self.annotations / name if self.annotations else None

        def load(p):
            rows = {}
            if p:
                for row in csv.reader(p.read_text(encoding="utf-8").splitlines()):
                    if row and row[0].lower().startswith("0x"):
                        rows[_int(row[0])] = row
            return rows
        return self._load(name, path, load)

    def _prototypes(self):
        path = self.annotations / "prototypes.h" if self.annotations else None

        def load(p):
            out = {}
            if p:
                lines = p.read_text(encoding="utf-8").split("\n")
                i = 0
                while i < len(lines):
                    line = lines[i]
                    m = re.match(r"^[A-Za-z_][^(;]*?\b([A-Za-z_]\w*)\s*\(", line) if not line.startswith(
                        ("typedef", "#", "/", " ")) else None
                    j = i + 1
                    while m and j < len(lines) and lines[j].startswith(" ") and not re.match(r"^[A-Za-z_]", lines[j]):
                        j += 1
                    if m:
                        out[m.group(1)] = "\n".join(lines[i:j])
                    i = j if m else i + 1
            return out
        return self._load("prototypes", path, load)

    # ----------------------------------------------------------------- names
    def function_at(self, addr):
        """The function containing addr: (start, size, name), from facts.jsonl, named from names.csv."""
        facts = self._facts()
        if addr in facts:
            start = addr
        else:
            starts = [a for a in facts if a <= addr < a + int(facts[a].get("size") or 0)]
            if not starts:
                return None
            start = max(starts)
        d = facts[start]
        return start, int(d.get("size") or 0), self.name_of(start) or d.get("name")

    def name_of(self, addr):
        row = self._csv("names.csv").get(addr)
        if row and len(row) > 1 and row[1]:
            return row[1]
        g = self._csv("globals.csv").get(addr)
        if g and len(g) > 1 and g[1]:
            return g[1]
        imp = self._imports().get(addr)
        if imp:
            return imp
        d = self._facts().get(addr)
        return d.get("name") if d else None

    def address_of(self, name):
        for table in ("names.csv", "globals.csv"):
            for addr, row in self._csv(table).items():
                if len(row) > 1 and row[1] == name:
                    return addr
        for addr, d in self._facts().items():
            if d.get("name") == name:
                return addr
        return None

    def resolve(self, target):
        """'0x4AD850', '4AD850', or a name -> address."""
        text = str(target).strip()
        if re.fullmatch(r"(0x)?[0-9A-Fa-f]{5,8}", text):
            return int(text, 16)
        addr = self.address_of(text)
        if addr is None:
            raise HarborError(f"{target!r} is neither an address nor a known name")
        return addr

    def convention(self, addr):
        """The calling convention read from the bytes: register inputs (a dataflow pass) and the `ret N` cleanup."""
        if not self.conventions_path or not self.conventions_path.exists():
            return None
        module = self._cache.get("conventions_module")
        if module is None:
            import importlib.util
            spec = importlib.util.spec_from_file_location("harbor_conventions", self.conventions_path)
            module = importlib.util.module_from_spec(spec)
            spec.loader.exec_module(module)
            self._cache["conventions_module"] = module
        func = self.function_at(addr)
        if not func:
            return None
        result = module.analyse(self._bytes(), func[0], func[1])
        if not result:
            return None
        return {"summary": module.describe(result), "registers_in": result["registers_in"],
                "ret_pops": result["ret_pops"]}

    # ----------------------------------------------------------------- tools
    def read(self, addr, count):
        data = self._bytes()
        offset = addr - self.base
        if offset < 0 or offset >= len(data):
            raise HarborError(f"0x{addr:X} is outside the image (0x{self.base:X}..0x{self.base + len(data):X})")
        return data[offset:offset + count]

    def disassemble(self, target, max_bytes=None):
        """The whole function (or max_bytes from an address) with call targets, jumps into other functions,
        import slots and known globals named."""
        try:
            import capstone
        except ImportError:
            raise HarborError("capstone is not installed for this Python (py -3.12 -m pip install capstone)") from None
        addr = self.resolve(target)
        func = self.function_at(addr)
        if max_bytes:
            start, size = addr, min(int(max_bytes), MAX_BYTES)
        elif func:
            start, size = func[0], func[1]
        else:
            start, size = addr, 0x100
        if size > MAX_BYTES:
            raise HarborError(f"the function is {size} bytes; pass max_bytes (at most {MAX_BYTES}) and walk it in parts")
        code = self.read(start, size)
        md = capstone.Cs(capstone.CS_ARCH_X86, capstone.CS_MODE_32)
        lines = []
        if func and not max_bytes:
            lines.append(f"; {func[2] or 'unnamed'} 0x{func[0]:X}..0x{func[0] + func[1]:X} ({func[1]} bytes)")
            conv = self.convention(func[0])
            if conv:
                lines.append(f"; convention from the bytes: {conv['summary']}; registers in: "
                             f"{', '.join(conv['registers_in']) or 'none'}; ret pops {conv['ret_pops']}")
        for ins in md.disasm(code, start):
            note = self._annotate(ins.op_str, ins.mnemonic, start, start + size)
            lines.append(f"{ins.address:08X}  {ins.mnemonic:7} {ins.op_str}{('    ; ' + note) if note else ''}")
        return "\n".join(lines)

    def _annotate(self, op_str, mnemonic, lo, hi):
        notes = []
        constant = re.search(r"(dword|qword) ptr \[0x([0-9a-f]{6,8})\]", op_str) if mnemonic.startswith("f") else None
        if constant:                                       # an x87 constant: show its value as the FPU reads it
            size = 4 if constant.group(1) == "dword" else 8
            try:
                raw = self.read(int(constant.group(2), 16), size)
                value = struct.unpack("<f" if size == 4 else "<d", raw)[0]
                notes.append(f"= {value!r}{'f' if size == 4 else ''}")
            except HarborError:
                pass
        for m in re.finditer(r"0x([0-9a-f]{6,8})", op_str):
            value = int(m.group(1), 16)
            if lo <= value < hi and not mnemonic.startswith("call"):
                continue                                   # a jump inside this function needs no name
            name = self.name_of(value)
            if name:
                notes.append(f"0x{value:X} {name}")
            elif mnemonic.startswith("call") or mnemonic.startswith("j"):
                func = self.function_at(value)
                if func and func[0] != value:
                    notes.append(f"0x{value:X} inside {func[2] or hex(func[0])}+0x{value - func[0]:X}")
        return ", ".join(notes)

    def xrefs(self, target):
        """Callers and callees of a function, and the data it touches, named."""
        addr = self.resolve(target)
        func = self.function_at(addr)
        if not func:
            raise HarborError(f"0x{addr:X} is not inside a known function")
        d = self._facts()[func[0]]
        named = lambda a: {"address": f"0x{a:X}", "name": self.name_of(a)}
        return {"function": named(func[0]) | {"size": func[1]},
                "callers": [named(a) for a in sorted(set(self._callers().get(func[0], [])))],
                "calls": [named(_int(a)) for a in d.get("calls", [])],
                "data": [named(_int(a)) for a in d.get("data", [])],
                "imports": d.get("imports", []), "strings": d.get("strings", [])}

    def bytes_at(self, target, count=16, as_type="hex"):
        """Raw bytes at an address, also read as hex, int32, uint32, float, double or a C string."""
        addr = self.resolve(target)
        count = max(1, min(int(count), 4096))
        raw = self.read(addr, count if as_type != "string" else 4096)
        if as_type == "string":
            text = raw.split(b"\0", 1)[0]
            return {"address": f"0x{addr:X}", "string": text.decode("latin-1"), "length": len(text)}
        if as_type in ("int32", "uint32", "float"):
            fmt = {"int32": "i", "uint32": "I", "float": "f"}[as_type]
            values = list(struct.unpack_from(f"<{len(raw) // 4}{fmt}", raw))
            return {"address": f"0x{addr:X}", as_type: [repr(v) if as_type == "float" else v for v in values],
                    "hex": raw[:len(values) * 4].hex(" ")}
        if as_type == "double":
            values = list(struct.unpack_from(f"<{len(raw) // 8}d", raw))
            return {"address": f"0x{addr:X}", "double": [repr(v) for v in values], "hex": raw[:len(values) * 8].hex(" ")}
        return {"address": f"0x{addr:X}", "hex": raw.hex(" "), "name": self.name_of(addr)}

    def lookup(self, target):
        """Everything the team has written about an address or name: names.csv, functions.csv and globals.csv rows,
        the prototype, the decompiled file and its status."""
        addr = self.resolve(target)
        out = {"address": f"0x{addr:X}"}
        func = self.function_at(addr)
        if func:
            out["function"] = {"start": f"0x{func[0]:X}", "size": func[1], "name": func[2]}
            conv = self.convention(func[0])
            if conv:
                out["convention_from_bytes"] = conv
        for table, key in (("names.csv", "names_row"), ("functions.csv", "functions_row"), ("globals.csv", "globals_row")):
            row = self._csv(table).get(addr)
            if row:
                out[key] = ",".join(row) if len(",".join(row)) < 4000 else ",".join(row)[:4000] + "..."
        name = self.name_of(addr)
        if name and name in self._prototypes():
            out["prototype"] = self._prototypes()[name]
        if self.src and self.src.exists() and func:
            prefix = f"{func[0]:06X}_"
            for path in self.src.rglob(f"{prefix}*.c"):
                head = path.read_text(encoding="utf-8", errors="replace").split("\n")[:6]
                status = next((l.split("status:", 1)[1].strip() for l in head if "status:" in l), "?")
                out["file"] = {"path": str(path.relative_to(self.src.parent)).replace(os.sep, "/"), "status": status}
                break
        return out


class Ghidra:
    """A small client for the shared, read-only Ghidra MCP server (GhidrAssistMCP over streamable HTTP), so every
    agent reaches Ghidra through Harbor's own tools. Configured by `ghidra_url` in the [reverse] table."""

    def __init__(self, config):
        self.url = (config or {}).get("ghidra_url")
        if not self.url:
            raise HarborError("this project has no ghidra_url in its [reverse] table")
        self.session = None
        self.counter = 0

    def _post(self, method, params=None, notify=False, timeout=120):
        import urllib.error
        import urllib.request
        body = {"jsonrpc": "2.0", "method": method, **({"params": params} if params is not None else {})}
        if not notify:
            self.counter += 1
            body["id"] = self.counter
        headers = {"Content-Type": "application/json", "Accept": "application/json, text/event-stream"}
        if self.session:
            headers["Mcp-Session-Id"] = self.session
        request = urllib.request.Request(self.url, json.dumps(body).encode(), headers=headers, method="POST")
        try:
            with urllib.request.urlopen(request, timeout=timeout) as response:
                self.session = response.headers.get("Mcp-Session-Id") or self.session
                text = response.read().decode("utf-8", "replace")
        except urllib.error.HTTPError as error:
            if error.code in (400, 404) and self.session and method != "initialize":
                self.session = None                       # the server restarted: start a new session once
                self._initialize()
                return self._post(method, params, notify, timeout)
            raise HarborError(f"the Ghidra server answered {error.code}: {error.read()[:300]!r}") from None
        except OSError as error:
            raise HarborError(f"the shared Ghidra server at {self.url} is not answering ({error}); start it with "
                              "R:/Workspace/Tools/re/ghidra-serve.ps1, or ask the coordinator") from None
        if notify or not text.strip():
            return None
        if not text.lstrip().startswith("{"):
            text = next((line[5:] for line in text.splitlines() if line.startswith("data:")), "{}")
        answer = json.loads(text)
        if "error" in answer:
            raise HarborError(f"Ghidra: {answer['error'].get('message', answer['error'])}")
        return answer.get("result")

    def _initialize(self):
        self._post("initialize", {"protocolVersion": "2025-06-18", "capabilities": {},
                                  "clientInfo": {"name": "harbor", "version": "2"}})
        self._post("notifications/initialized", notify=True)

    def _ensure(self):
        if not self.session:
            self._initialize()

    def tools(self):
        """The server's tools: name, description and argument schema."""
        self._ensure()
        return [{"name": t["name"], "description": t.get("description", ""),
                 "arguments": (t.get("inputSchema") or {}).get("properties", {}),
                 "required": (t.get("inputSchema") or {}).get("required", [])}
                for t in self._post("tools/list")["tools"]]

    def call(self, tool, arguments=None, wait=180):
        """Call one tool; long-running ones (the decompiler) are polled until they finish or `wait` seconds pass."""
        import time
        self._ensure()
        result = self._post("tools/call", {"name": tool, "arguments": arguments or {}}, timeout=wait + 30)
        text = "\n".join(c.get("text", "") for c in result.get("content", []) if isinstance(c, dict))
        task = re.search(r"Task ID: ([0-9a-f-]{36})", text) if "async execution" in text else None
        deadline = time.monotonic() + wait
        while task and time.monotonic() < deadline:
            time.sleep(1)
            result = self._post("tools/call", {"name": "get_task_status", "arguments": {"task_id": task.group(1)}})
            text = "\n".join(c.get("text", "") for c in result.get("content", []) if isinstance(c, dict))
            if not re.search(r"Status: (RUNNING|PENDING|QUEUED)", text):
                break
        if result.get("isError"):
            raise HarborError(text.strip() or f"Ghidra refused {tool}")
        return re.sub(r"^\[Context\] Operating on: [^\n]*\n+", "", text)


class Debugger:
    """A per-agent x64dbg session through the x64dbg-automate MCP server (stdio), started on first use and kept for
    the life of this Harbor process, so breakpoints and the paused game survive between calls. Configured by
    `debugger_command`, `debugger_path` (x32dbg.exe) and `debugger_temp` in the [reverse] table. The plugin and the
    client find each other through a session file in TEMP, so both get the same TEMP."""

    def __init__(self, config):
        config = config or {}
        self.command = config.get("debugger_command")
        if not self.command:
            raise HarborError("this project has no debugger_command in its [reverse] table")
        self.env = dict(os.environ)
        if config.get("debugger_path"):
            self.env["X64DBG_PATH"] = str(config["debugger_path"])
        if config.get("debugger_temp"):
            self.env["TEMP"] = self.env["TMP"] = str(config["debugger_temp"])
        self.proc = None
        self.counter = 0
        self.lock = __import__("threading").Lock()

    def _start(self):
        import atexit
        import queue
        import subprocess
        import threading
        command = self.command if isinstance(self.command, list) else [self.command]
        self.proc = subprocess.Popen(command, stdin=subprocess.PIPE, stdout=subprocess.PIPE,
                                     stderr=subprocess.DEVNULL, env=self.env, text=True, encoding="utf-8",
                                     creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        self.answers = queue.Queue()

        def read(stream, answers):
            for line in stream:
                if line.strip():
                    try:
                        answers.put(json.loads(line))
                    except ValueError:
                        pass
            answers.put(None)
        threading.Thread(target=read, args=(self.proc.stdout, self.answers), daemon=True).start()
        self._rpc("initialize", {"protocolVersion": "2025-06-18", "capabilities": {},
                                 "clientInfo": {"name": "harbor", "version": "2"}})
        self._send({"jsonrpc": "2.0", "method": "notifications/initialized"})
        atexit.register(self.close)

    def _send(self, body):
        self.proc.stdin.write(json.dumps(body) + "\n")
        self.proc.stdin.flush()

    def _rpc(self, method, params, timeout=120):
        import queue
        self.counter += 1
        self._send({"jsonrpc": "2.0", "id": self.counter, "method": method, "params": params})
        while True:
            try:
                answer = self.answers.get(timeout=timeout)
            except queue.Empty:
                raise HarborError(f"the debugger did not answer {method} within {timeout}s") from None
            if answer is None:
                self.proc = None
                raise HarborError("the debugger's MCP server exited")
            if answer.get("id") == self.counter:
                if "error" in answer:
                    raise HarborError(f"debugger: {answer['error'].get('message', answer['error'])}")
                return answer.get("result")

    def tools(self):
        with self.lock:
            if self.proc is None or self.proc.poll() is not None:
                self._start()
            return [{"name": t["name"], "description": t.get("description", ""),
                     "arguments": (t.get("inputSchema") or {}).get("properties", {}),
                     "required": (t.get("inputSchema") or {}).get("required", [])}
                    for t in self._rpc("tools/list", {})["tools"]]

    def call(self, tool, arguments=None, timeout=300):
        with self.lock:
            if self.proc is None or self.proc.poll() is not None:
                self._start()
            result = self._rpc("tools/call", {"name": tool, "arguments": arguments or {}}, timeout=timeout)
        text = "\n".join(c.get("text", "") for c in result.get("content", []) if isinstance(c, dict))
        if result.get("isError"):
            raise HarborError(text.strip() or f"the debugger refused {tool}")
        return text

    def close(self):
        """End the debuggee and the debugger with the Harbor process, so no paused game is left behind."""
        if self.proc and self.proc.poll() is None:
            try:
                self._rpc("tools/call", {"name": "terminate_session", "arguments": {}}, timeout=15)
            except HarborError:
                pass
            self.proc.terminate()
        self.proc = None


def match_functions(config, targets, decomp=None, flags=None):
    """Compile functions with the game's compiler (Visual C++ 2003) and compare them with the original bytes.
    targets: addresses ("0x4ADC70, 0x412F70"); decomp: the checkout or worktree whose src/ holds them."""
    import subprocess
    import sys
    import tempfile
    config = config or {}
    matcher = config.get("matcher")
    if not matcher:
        raise HarborError("this project has no matcher in its [reverse] table")
    wanted = [t.strip() for t in re.split(r"[,\s]+", str(targets)) if t.strip()]
    if not wanted or len(wanted) > 50:
        raise HarborError("give 1 to 50 addresses")
    with tempfile.TemporaryDirectory(dir=config.get("debugger_temp")) as tmp:
        out = Path(tmp) / "match.json"
        command = [sys.executable, "-B", str(matcher), "--all", "--only", ",".join(wanted), "--json", str(out),
                   "--flags", flags or config.get("match_flags", "-TC -O2")]
        if decomp:
            command += ["--decomp", str(decomp)]
        result = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", errors="replace",
                                timeout=600, creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        if not out.exists():
            raise HarborError("the matcher failed: " + (result.stdout + result.stderr).strip()[-800:])
        rows = json.loads(out.read_text(encoding="utf-8"))
    errors = [line for line in result.stdout.splitlines() if "compile errors" in line]
    return {"results": rows, "summary": result.stdout.strip().splitlines()[0] if result.stdout.strip() else "",
            **({"compile": errors[0]} if errors else {})}


def run_toolkit(config, key, argv, json_to_file=False, timeout=900):
    """Run the toolkit script configured as [reverse] <key> with argv and return its JSON (stdout, or a temp file
    passed as `--json <file>` when json_to_file)."""
    import subprocess
    import sys
    import tempfile
    script = (config or {}).get(key)
    if not script or not Path(script).exists():
        raise HarborError(f"this project has no {key} script in its [reverse] table")
    with tempfile.TemporaryDirectory(dir=(config or {}).get("debugger_temp")) as tmp:
        out = Path(tmp) / "out.json"
        command = [sys.executable, "-B", str(script), *argv] + (["--json", str(out)] if json_to_file else ["--json"])
        result = subprocess.run(command, capture_output=True, text=True, encoding="utf-8", errors="replace",
                                timeout=timeout, cwd=str(Path(script).parent),
                                creationflags=getattr(subprocess, "CREATE_NO_WINDOW", 0))
        text = out.read_text(encoding="utf-8") if json_to_file and out.exists() else result.stdout
        try:
            return json.loads(text)
        except ValueError:
            raise HarborError(f"{key} failed: " + (result.stdout + result.stderr).strip()[-800:]) from None