from __future__ import annotations

from typing import TYPE_CHECKING, Any

from aicd.tools.script_backup import policy_allow_under_ai_output_or_workspace

from aicd.agent.tool_args_coerce import (
    coerce_tool_json_array,
    coerce_tool_json_object,
)
from aicd.tools.process_tool import (
    shell_argv_from_user_input,
    warn_shell_preflight,
)

if TYPE_CHECKING:
    from aicd.agent.tool_router import ToolRouter


def _text_regex_replace(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    return router._text.regex_replace(
        args["path"],
        args["pattern"],
        args["replacement"],
        count=int(args.get("count") or 0),
    )


def _text_patch_bytes_hex(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    return router._text.patch_bytes(
        args["path"], int(args["offset"]), args["hex_data"]
    )


async def _run_command(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    timeout = float(args.get("timeout_sec", 600))
    raw_argv = args.get("argv")
    shell = args.get("shell")
    if raw_argv is not None:
        co = coerce_tool_json_array(raw_argv, "argv", allow_empty=True)
        if isinstance(co, dict):
            return co
        raw_argv = co
    cmd = args.get("command")
    if raw_argv is None and cmd is not None and str(cmd).strip():
        if (w := warn_shell_preflight(str(cmd))):
            return {"ok": False, "error": w}
    try:
        argv = shell_argv_from_user_input(
            command=args.get("command"),
            argv=[str(x) for x in raw_argv] if raw_argv else None,
            shell=str(shell).strip() if shell is not None else None,
        )
    except ValueError as e:
        return {"ok": False, "error": str(e)}

    async def _log(line: str) -> None:
        await router._tasks.bus.publish(
            {"topic": "agent", "event": "log", "line": line}
        )

    return await router._process.run_command(
        argv,
        cwd=args.get("cwd"),
        timeout_sec=timeout,
        on_line=_log,
    )


async def _http_request(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    hdrs = args.get("headers")
    if hdrs is not None:
        h_obj, h_err = coerce_tool_json_object(hdrs, "headers")
        if h_err:
            return {"ok": False, "error": h_err}
        hdrs = h_obj
    jb = args.get("json_body")
    if jb is not None:
        j_obj, j_err = coerce_tool_json_object(jb, "json_body")
        if j_err:
            return {"ok": False, "error": j_err}
        jb = j_obj
    return await router._http.request(
        args["method"],
        args["url"],
        headers=hdrs,
        json_body=jb,
        data=args.get("data"),
        timeout_sec=float(args.get("timeout_sec", 60)),
    )


async def _http_download(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    return await router._http.download(
        args["url"], filename=args.get("filename")
    )


async def _script_run(router: ToolRouter, args: dict[str, Any]) -> dict[str, Any]:
    kind = str(args["kind"]).lower()
    src = args["source"]
    timeout = float(args.get("timeout_sec", 300))
    rd = args.get("related_doc")
    related_doc = str(rd).strip() if isinstance(rd, str) and rd.strip() else None
    bak_root = (router._ai_output / "bak").resolve()
    bak_policy = policy_allow_under_ai_output_or_workspace(
        router._ai_output, router._policy.check_allowed
    )

    async def _log(line: str) -> None:
        await router._tasks.bus.publish(
            {"topic": "agent", "event": "log", "line": line}
        )

    common = dict(
        on_line=_log,
        timeout_sec=timeout,
        related_doc=related_doc,
        backup_dir=bak_root,
        policy_check=bak_policy,
    )
    if kind == "py":
        return await router._scripts.run_python(src, **common)
    if kind == "ps1":
        return await router._scripts.run_powershell(src, **common)
    if kind == "cmd":
        return await router._scripts.run_cmd(src, **common)
    if kind == "bash":
        return await router._scripts.run_bash(src, **common)
    return {"ok": False, "error": f"unknown script kind: {kind}"}


HANDLERS: dict[str, Any] = {
    "text_regex_replace": _text_regex_replace,
    "text_patch_bytes_hex": _text_patch_bytes_hex,
    "run_command": _run_command,
    "http_request": _http_request,
    "http_download": _http_download,
    "script_run": _script_run,
}
