"""Tool registry and dispatch for the security sidecar.

Each tool is a callable(subprocess_args, timeout_ms) -> ToolResult.
The registry maps tool names to implementations and provides health checking.
"""

import json
import os
import shlex
import shutil
import signal
import subprocess
import sys
import tempfile
import threading
from dataclasses import dataclass, field
from datetime import datetime, timezone
from pathlib import Path
from typing import Optional


@dataclass
class ToolResult:
    ok: bool
    output: str = ""
    error: str = ""
    returncode: int = -1
    timed_out: bool = False
    duration_ms: int = 0


def _run_subprocess(
    cmd: list[str],
    stdin_data: Optional[bytes] = None,
    timeout_ms: int = 30_000,
    cwd: Optional[Path] = None,
    env: Optional[dict[str, str]] = None,
) -> ToolResult:
    start = datetime.now(timezone.utc)
    try:
        proc = subprocess.Popen(
            cmd,
            stdin=subprocess.PIPE,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            cwd=cwd,
            env=env,
            preexec_fn=lambda: signal.signal(signal.SIGXCPU, signal.SIG_DFL),
        )
        stdout_b, stderr_b = b"", b""
        done = threading.Event()

        def _reader():
            nonlocal stdout_b, stderr_b
            try:
                stdout_b, stderr_b = proc.communicate(input=stdin_data, timeout=timeout_ms / 1000)
            except subprocess.TimeoutExpired:
                proc.kill()
                stdout_b, stderr_b = proc.communicate()
            finally:
                done.set()

        reader_thread = threading.Thread(target=_reader, daemon=True)
        reader_thread.start()
        reader_thread.join(timeout=(timeout_ms / 1000) + 2)
        if not done.is_set():
            proc.kill()
            reader_thread.join(1)
            elapsed = int((datetime.now(timezone.utc) - start).total_seconds() * 1000)
            return ToolResult(
                ok=False,
                error=f"timed out after {timeout_ms}ms",
                returncode=-signal.SIGKILL,
                timed_out=True,
                duration_ms=elapsed,
            )

        combined = stdout_b.decode("utf-8", errors="replace")
        if stderr_b:
            combined += "\n" + stderr_b.decode("utf-8", errors="replace")
        elapsed = int((datetime.now(timezone.utc) - start).total_seconds() * 1000)
        return ToolResult(
            ok=proc.returncode == 0,
            output=combined if proc.returncode == 0 else "",
            error=combined if proc.returncode != 0 else "",
            returncode=proc.returncode or 0,
            duration_ms=elapsed,
        )
    except FileNotFoundError:
        elapsed = int((datetime.now(timezone.utc) - start).total_seconds() * 1000)
        return ToolResult(
            ok=False,
            error=f"executable not found: {cmd[0]}",
            duration_ms=elapsed,
        )
    except Exception as e:
        elapsed = int((datetime.now(timezone.utc) - start).total_seconds() * 1000)
        return ToolResult(ok=False, error=str(e), duration_ms=elapsed)


def _check_tool(name: str) -> bool:
    return shutil.which(name) is not None


# ── Web tools ──────────────────────────────────────────────────────

def _run_http(args: dict, timeout_ms: int) -> ToolResult:
    url = args.get("url", "")
    method = args.get("method", "GET").upper()
    headers = args.get("headers", {})
    data = args.get("data", "")
    if not url:
        return ToolResult(ok=False, error="url is required")
    cmd = ["curl", "-s", "-S", "-L", "-X", method]
    for k, v in headers.items():
        cmd.extend(["-H", f"{k}: {v}"])
    if data and method in ("POST", "PUT", "PATCH"):
        cmd.extend(["-d", data])
    cmd.append(url)
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_sqlmap(args: dict, timeout_ms: int) -> ToolResult:
    url = args.get("url", "")
    if not url:
        return ToolResult(ok=False, error="url is required")
    cmd = ["sqlmap", "--batch", "--random-agent", "--time-sec", "5"]
    if args.get("cookie"):
        cmd.extend(["--cookie", args["cookie"]])
    if args.get("data"):
        cmd.extend(["--data", args["data"]])
    if args.get("level"):
        cmd.extend(["--level", str(args["level"])])
    if args.get("risk"):
        cmd.extend(["--risk", str(args["risk"])])
    cmd.append(url)
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_nuclei(args: dict, timeout_ms: int) -> ToolResult:
    target = args.get("target", "")
    if not target:
        return ToolResult(ok=False, error="target is required")
    cmd = ["nuclei", "-silent", "-no-color"]
    if args.get("templates"):
        cmd.extend(["-t", args["templates"]])
    if args.get("severity"):
        cmd.extend(["-severity", args["severity"]])
    cmd.extend(["-u", target])
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_ffuf(args: dict, timeout_ms: int) -> ToolResult:
    url = args.get("url", "")
    wordlist = args.get("wordlist", "/usr/share/wordlists/dirb/common.txt")
    if not url:
        return ToolResult(ok=False, error="url is required")
    cmd = ["ffuf", "-u", url, "-w", wordlist, "-ac", "-t", "40"]
    if args.get("extensions"):
        cmd.extend(["-e", args["extensions"]])
    if args.get("fc"):
        cmd.extend(["-fc", str(args["fc"])])
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_dalfox(args: dict, timeout_ms: int) -> ToolResult:
    url = args.get("url", "")
    if not url:
        return ToolResult(ok=False, error="url is required")
    cmd = ["dalfox", "url", url, "--silence", "--no-color", "--only-poc", "gfm"]
    if args.get("cookie"):
        cmd.extend(["--cookie", args["cookie"]])
    if args.get("param"):
        cmd.extend(["-p", args["param"]])
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_zap(args: dict, timeout_ms: int) -> ToolResult:
    target = args.get("target", "")
    if not target:
        return ToolResult(ok=False, error="target is required")
    cmd = ["zap-cli", "--silent", "quick-scan", "-t", str(args.get("timeout", 60))]
    if args.get("spider"):
        cmd.append("--spider")
    cmd.append(target)
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_xss_confirm(args: dict, timeout_ms: int) -> ToolResult:
    payloads = args.get("payloads", [
        "<script>alert(1)</script>",
        "\"><script>alert(1)</script>",
        "';alert(1)//",
    ])
    url_template = args.get("url", "")
    param = args.get("param", "q")
    if not url_template:
        return ToolResult(ok=False, error="url template with {payload} placeholder is required")
    for payload in payloads:
        url = url_template.replace("{payload}", payload)
        try:
            resp = _run_subprocess(
                ["curl", "-s", "-o", "/dev/null", "-w", "%{http_code}", url],
                timeout_ms=timeout_ms // len(payloads),
            )
            if resp.ok and resp.output.strip() not in ("404", "400"):
                return ToolResult(
                    ok=True,
                    output=f"XSS vector confirmed: {payload} returned HTTP {resp.output.strip()}",
                )
        except Exception:
            continue
    return ToolResult(ok=False, error="no XSS vectors confirmed from payload set")


# ── Crypto tools ───────────────────────────────────────────────────

def _run_z3(args: dict, timeout_ms: int) -> ToolResult:
    script = args.get("script", "")
    if not script:
        return ToolResult(ok=False, error="z3 script (SMT-LIB or Python) is required")
    ext = ".smt2" if script.strip().startswith("(") else ".py"
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=ext, delete=False, prefix="z3_"
    ) as f:
        f.write(script)
        tmp = f.name
    try:
        if ext == ".py":
            cmd = ["python3", tmp]
        else:
            cmd = ["z3", "-in", tmp]
        result = _run_subprocess(cmd, timeout_ms=timeout_ms)
        os.unlink(tmp)
        return result
    except Exception:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def _run_sage(args: dict, timeout_ms: int) -> ToolResult:
    script = args.get("script", "")
    if not script:
        return ToolResult(ok=False, error="sage script is required")
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".sage", delete=False, prefix="sage_"
    ) as f:
        f.write(script)
        tmp = f.name
    try:
        result = _run_subprocess(["sage", tmp], timeout_ms=timeout_ms)
        os.unlink(tmp)
        return result
    except Exception:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def _run_rsa_tool(args: dict, timeout_ms: int) -> ToolResult:
    n = args.get("n", "")
    e = args.get("e", "65537")
    if not n:
        return ToolResult(ok=False, error="n (modulus) is required")
    script = f"""
import sys
try:
    from factordb.factordb import FactorDB
    n = int("{n}")
    f = FactorDB(n)
    f.connect()
    factors = f.get_factor_list()
    if factors:
        print("factors:", factors)
    else:
        print("no factors found from FactorDB")
except ImportError:
    print("factordb-python not available, trying z3...")
except Exception as e:
    print(f"error: {{e}}")
"""
    return _run_subprocess(
        ["python3", "-c", script.replace("{", "{{").replace("}", "}}")],
        timeout_ms=timeout_ms,
    )


def _run_factordb(args: dict, timeout_ms: int) -> ToolResult:
    n = args.get("n", "")
    if not n:
        return ToolResult(ok=False, error="n is required")
    script = f"""
import sys, json
try:
    from factordb.factordb import FactorDB
    f = FactorDB({n})
    f.connect()
    factors = f.get_factor_list()
    print(json.dumps({{"factors": factors, "status": f.get_status()}}))
except ImportError:
    import urllib.request
    url = f"http://factordb.com/api?query={n}"
    resp = urllib.request.urlopen(url, timeout=10)
    print(resp.read().decode())
except Exception as e:
    print(json.dumps({{"error": str(e)}}))
"""
    return _run_subprocess(["python3", "-c", script], timeout_ms=timeout_ms)


def _run_hashcat(args: dict, timeout_ms: int) -> ToolResult:
    hash_value = args.get("hash", "")
    mode = args.get("mode", "0")
    wordlist = args.get("wordlist", "/usr/share/wordlists/rockyou.txt")
    if not hash_value:
        return ToolResult(ok=False, error="hash is required")
    cmd = ["hashcat", "--force", "-m", mode, "-a", "0", hash_value, wordlist]
    if args.get("rules"):
        cmd.extend(["-r", args["rules"]])
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_hashid(args: dict, timeout_ms: int) -> ToolResult:
    hash_value = args.get("hash", "")
    if not hash_value:
        return ToolResult(ok=False, error="hash is required")
    cmd = ["hashid", "-m", hash_value]
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_decode(args: dict, timeout_ms: int) -> ToolResult:
    encoded = args.get("data", "")
    encoding = args.get("encoding", "base64")
    if not encoded:
        return ToolResult(ok=False, error="data is required")
    script = f"""
import base64, binascii, json
data = {json.dumps(encoded)}
enc = {json.dumps(encoding)}
try:
    if enc == "base64":
        result = base64.b64decode(data).decode("utf-8", errors="replace")
    elif enc == "base32":
        result = base64.b32decode(data).decode("utf-8", errors="replace")
    elif enc == "hex":
        result = bytes.fromhex(data).decode("utf-8", errors="replace")
    elif enc == "rot13":
        import codecs
        result = codecs.decode(data, "rot_13")
    else:
        result = f"unknown encoding: {{enc}}"
    print(result)
except Exception as e:
    print(f"decode failed: {{e}}")
"""
    return _run_subprocess(["python3", "-c", script], timeout_ms=timeout_ms)


# ── Reverse-engineering tools ──────────────────────────────────────

def _run_js_deobf(args: dict, timeout_ms: int) -> ToolResult:
    source = args.get("source", "")
    if not source:
        return ToolResult(ok=False, error="source is required")
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".js", delete=False, prefix="jsdeob_"
    ) as f:
        f.write(source)
        tmp = f.name
    try:
        cmd = ["npx", "--yes", "deobfuscate-js", tmp]
        result = _run_subprocess(cmd, timeout_ms=timeout_ms)
        os.unlink(tmp)
        return result
    except Exception:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def _run_sourcemap(args: dict, timeout_ms: int) -> ToolResult:
    url = args.get("url", "")
    if not url:
        return ToolResult(ok=False, error="url is required")
    cmd = ["curl", "-s", "-L", url]
    result = _run_subprocess(cmd, timeout_ms=timeout_ms)
    if not result.ok:
        return result
    try:
        import base64
        import json
        data = json.loads(result.output)
        if "mappings" in data:
            sources = data.get("sources", [])
            return ToolResult(ok=True, output=json.dumps({"sources": sources, "file": data.get("file", "")}))
        return ToolResult(ok=False, error="response is not a valid source map")
    except json.JSONDecodeError as e:
        return ToolResult(ok=False, error=f"invalid JSON: {e}")


def _run_wasm_decompile(args: dict, timeout_ms: int) -> ToolResult:
    wasm_path = args.get("path", "")
    wasm_data = args.get("data", "")
    if wasm_path:
        cmd = ["wasm-decompile", wasm_path]
    elif wasm_data:
        with tempfile.NamedTemporaryFile(
            mode="wb", suffix=".wasm", delete=False, prefix="wasm_"
        ) as f:
            import base64
            f.write(base64.b64decode(wasm_data))
            tmp = f.name
        try:
            cmd = ["wasm-decompile", tmp]
            result = _run_subprocess(cmd, timeout_ms=timeout_ms)
            os.unlink(tmp)
            return result
        except Exception:
            if os.path.exists(tmp):
                os.unlink(tmp)
            raise
    else:
        return ToolResult(ok=False, error="path or base64-encoded data is required")
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


# ── Pwn tools ──────────────────────────────────────────────────────

def _run_triage(args: dict, timeout_ms: int) -> ToolResult:
    binary = args.get("binary", "")
    if not binary:
        return ToolResult(ok=False, error="binary path is required")
    cmd = ["file", binary]
    file_result = _run_subprocess(cmd, timeout_ms=5000)
    checksec_cmd = ["checksec", "--file=" + binary]
    check_result = _run_subprocess(checksec_cmd, timeout_ms=5000)
    combined = file_result.output or ""
    if check_result.output:
        combined += "\n" + check_result.output
    return ToolResult(ok=True, output=combined)


def _run_ropgadget(args: dict, timeout_ms: int) -> ToolResult:
    binary = args.get("binary", "")
    if not binary:
        return ToolResult(ok=False, error="binary path is required")
    cmd = ["ROPgadget", "--binary", binary]
    if args.get("depth"):
        cmd.extend(["--depth", str(args["depth"])])
    if args.get("only"):
        cmd.extend(["--only", args["only"]])
    if args.get("range"):
        cmd.extend(["--range", args["range"]])
    cmd.append("--silent")
    return _run_subprocess(cmd, timeout_ms=timeout_ms)


def _run_pwntools(args: dict, timeout_ms: int) -> ToolResult:
    script = args.get("script", "")
    if not script:
        return ToolResult(ok=False, error="pwntools Python script is required")
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".py", delete=False, prefix="pwn_"
    ) as f:
        f.write("#!/usr/bin/env python3\n")
        f.write("from pwn import *\n")
        f.write("context.log_level = 'error'\n")
        f.write(script)
        tmp = f.name
    try:
        result = _run_subprocess(["python3", tmp], timeout_ms=timeout_ms)
        os.unlink(tmp)
        return result
    except Exception:
        if os.path.exists(tmp):
            os.unlink(tmp)
        raise


def _run_exploit_template(args: dict, timeout_ms: int) -> ToolResult:
    target = args.get("target", "")
    template_type = args.get("type", "ret2libc")
    if not target:
        return ToolResult(ok=False, error="target binary path is required")
    arch = args.get("arch", "amd64")
    template = f"""#!/usr/bin/env python3
from pwn import *
context.binary = '{target}'
context.arch = '{arch}'
context.log_level = 'warn'

elf = ELF('{target}')
"""
    if template_type == "ret2libc":
        template += f"""
# ret2libc template
rop = ROP(elf)
pop_rdi = rop.find_gadget(['pop rdi', 'ret'])
if pop_rdi:
    pop_rdi = pop_rdi[0]
    print(f"pop rdi; ret @ {{hex(pop_rdi)}}")
bin_sh = next(elf.search(b'/bin/sh'), None)
if bin_sh:
    print(f"/bin/sh @ {{hex(bin_sh)}}")
system = elf.plt.get('system')
if system:
    print(f"system @ {{hex(system)}}")
else:
    libc = elf.libc
    if libc:
        print(f"libc: {{libc.path}}")
"""
    elif template_type == "shellcode":
        template += f"""
# shellcode execution template
shellcode = asm(shellcraft.sh())
print(f"shellcode ({len(shellcode)} bytes): {{shellcode.hex()}}")
"""
    else:
        template += f"""
# {template_type} exploit template
print(f"Target: {{elf.path}}")
print(f"PIE: {{elf.pie}}")
print(f"NX: {{elf.nx}}")
print(f"Canary: {{elf.canary}}")
"""
    return ToolResult(ok=True, output=template)


# ── Tool registry ──────────────────────────────────────────────────

class ToolRegistry:
    def __init__(self):
        self._tools: dict[str, callable] = {
            "http": _run_http,
            "sqlmap": _run_sqlmap,
            "nuclei": _run_nuclei,
            "ffuf": _run_ffuf,
            "dalfox": _run_dalfox,
            "zap": _run_zap,
            "xss_confirm": _run_xss_confirm,
            "z3": _run_z3,
            "sage": _run_sage,
            "rsa": _run_rsa_tool,
            "factordb": _run_factordb,
            "hashcat": _run_hashcat,
            "hashid": _run_hashid,
            "decode": _run_decode,
            "js_deobfuscate": _run_js_deobf,
            "sourcemap": _run_sourcemap,
            "wasm_decompile": _run_wasm_decompile,
            "triage": _run_triage,
            "ropgadget": _run_ropgadget,
            "pwntools": _run_pwntools,
            "exploit_template": _run_exploit_template,
        }

    def list_tools(self) -> list[str]:
        return list(self._tools.keys())

    def has_tool(self, name: str) -> bool:
        return name in self._tools

    def run(self, name: str, args: dict, timeout_ms: int = 30_000) -> ToolResult:
        if name not in self._tools:
            return ToolResult(ok=False, error=f"unknown tool: {name}")
        fn = self._tools[name]
        return fn(args, timeout_ms)

    def health_check(self) -> dict:
        results = {}
        for name in self._tools:
            binary = _binary_for_tool(name)
            if binary:
                results[name] = {"available": _check_tool(binary)}
            else:
                results[name] = {"available": True}  # Python-based, assume available
        return {
            "tools": results,
            "available_count": sum(1 for v in results.values() if v["available"]),
            "total_count": len(results),
        }


def _binary_for_tool(name: str) -> Optional[str]:
    mapping = {
        "http": "curl",
        "sqlmap": "sqlmap",
        "nuclei": "nuclei",
        "ffuf": "ffuf",
        "dalfox": "dalfox",
        "zap": "zap-cli",
        "xss_confirm": "curl",
        "z3": "z3",
        "sage": "sage",
        "hashcat": "hashcat",
        "hashid": "hashid",
        "pwntools": "python3",
        "ropgadget": "ROPgadget",
    }
    return mapping.get(name)