"""LangChain tools the inner deepagents loop uses.

  - workspace file CRUD (boto3 → MinIO, scoped to the user's bucket)
  - sandbox python (for `a2a card` / `a2a validate` round-trips)
  - cp_deploy_tarball (call /v1/agents/from-tarball on the user's behalf)
"""
from __future__ import annotations

import asyncio
import base64
import io
import json
import re
import tarfile
import time
from dataclasses import dataclass
from importlib import resources
from typing import TYPE_CHECKING, Any

import boto3
import httpx
from botocore.config import Config as _BotoConfig
from botocore.exceptions import ClientError
from langchain_core.tools import tool

if TYPE_CHECKING:
    from .config import Settings


@dataclass(frozen=True)
class ToolContext:
    bucket: str
    settings: "Settings"
    cp_jwt: str | None = None


async def _wait_for_live_card(
    url: str | None,
    expected_version: str | None,
    timeout_s: float = 180.0,
) -> tuple[bool, dict[str, Any]]:
    """Poll ``{url}/.well-known/agent-card`` until ``version`` matches
    ``expected_version`` or ``timeout_s`` elapses. Returns
    ``(matched, last_seen_card)``. ``last_seen_card`` is ``{}`` if the
    endpoint never responded with a parseable card.

    Backoff: 2s → 1.5x per try, capped at 10s. Tolerates pod-not-ready
    (any HTTPError or 4xx/5xx) — just keeps polling.
    """
    if not url or not expected_version:
        return False, {}
    deadline = time.monotonic() + timeout_s
    last_card: dict[str, Any] = {}
    interval = 2.0
    while time.monotonic() < deadline:
        try:
            async with httpx.AsyncClient(timeout=5.0) as c:
                r = await c.get(f"{url.rstrip('/')}/.well-known/agent-card")
            if r.status_code == 200:
                card = r.json() or {}
                if isinstance(card, dict):
                    last_card = card
                    if card.get("version") == expected_version:
                        return True, card
        except (httpx.HTTPError, ValueError):
            pass
        await asyncio.sleep(interval)
        interval = min(interval * 1.5, 10.0)
    return False, last_card


def _s3(ctx: ToolContext) -> Any:
    return boto3.client(
        "s3",
        endpoint_url=ctx.settings.minio_endpoint,
        aws_access_key_id=ctx.settings.minio_access_key,
        aws_secret_access_key=ctx.settings.minio_secret_key,
        region_name="us-east-1",
        config=_BotoConfig(s3={"addressing_style": "path"}),
    )


def _ensure_bucket(s3: Any, bucket: str) -> None:
    try:
        s3.head_bucket(Bucket=bucket)
    except ClientError as exc:
        code = exc.response.get("Error", {}).get("Code")
        if code in {"404", "NoSuchBucket", "NotFound"}:
            s3.create_bucket(Bucket=bucket)
        elif code != "403":
            raise


_AGENT_NAME_RE = re.compile(r"^[a-z][a-z0-9-]{1,62}$")


def _validate_name(name: str) -> None:
    if not _AGENT_NAME_RE.match(name):
        raise ValueError(
            f"invalid agent name {name!r}: must match {_AGENT_NAME_RE.pattern}"
        )


def _agent_prefix(name: str) -> str:
    _validate_name(name)
    return f"agents/{name}/"


def build_tools(ctx: ToolContext) -> list[Any]:
    bucket = ctx.bucket
    settings = ctx.settings
    s3 = _s3(ctx)
    _ensure_bucket(s3, bucket)

    @tool
    def init_agent_template(
        name: str,
        description: str = "A new A2A agent",
    ) -> str:
        """Initialize ``agents/<name>/`` from the installed a2a-pack template.

        Use this FIRST for a new project. It writes the same baseline files
        as ``a2a init`` for the SDK version installed in this runtime:
        ``agent.py``, ``a2a.yaml``, and ``requirements.txt``. After this,
        edit those files with ``read_agent_file`` + ``write_agent_file`` to
        implement the user's requested behavior.
        """
        try:
            files = _render_a2a_init_template(name, description=description)
            prefix = _agent_prefix(name)
        except (OSError, RuntimeError, ValueError) as exc:
            return json.dumps({"error": str(exc)})
        written: list[dict[str, Any]] = []
        for path, content in files.items():
            key = prefix + path
            s3.put_object(
                Bucket=bucket,
                Key=key,
                Body=content.encode("utf-8"),
                ContentType="text/plain; charset=utf-8",
            )
            written.append({"path": key, "size": len(content)})
        return json.dumps({"ok": True, "agent": name, "files": written})

    @tool
    def list_agent_files(name: str) -> str:
        """List every file under ``agents/<name>/`` in the user's workspace.

        Use this before edits to see what you already wrote / what's left.
        """
        try:
            prefix = _agent_prefix(name)
        except ValueError as exc:
            return json.dumps({"error": str(exc)})
        out: list[dict[str, Any]] = []
        paginator = s3.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page.get("Contents") or ():
                out.append({
                    "path": obj["Key"][len(prefix):],
                    "size": int(obj.get("Size") or 0),
                })
        return json.dumps({"agent": name, "files": out})

    @tool
    def write_agent_file(name: str, path: str, content: str) -> str:
        """Write a file under ``agents/<name>/<path>`` in the user's workspace.

        Use this for ``agent.py``, ``a2a.yaml``, ``requirements.txt``, and
        any auxiliary modules the agent needs.
        """
        try:
            prefix = _agent_prefix(name)
        except ValueError as exc:
            return json.dumps({"error": str(exc)})
        key = prefix + path.lstrip("/")
        s3.put_object(
            Bucket=bucket, Key=key, Body=content.encode("utf-8"),
            ContentType="text/plain; charset=utf-8",
        )
        return json.dumps({"ok": True, "path": key, "size": len(content)})

    @tool
    def read_agent_file(name: str, path: str) -> str:
        """Read a single file from the scaffolded agent project."""
        try:
            prefix = _agent_prefix(name)
        except ValueError as exc:
            return json.dumps({"error": str(exc)})
        key = prefix + path.lstrip("/")
        try:
            obj = s3.get_object(Bucket=bucket, Key=key)
        except ClientError as exc:
            return json.dumps({"error": f"read failed: {exc}"})
        body = obj["Body"].read().decode("utf-8", errors="replace")
        return json.dumps({"path": key, "content": body[:200_000]})

    @tool
    async def test_agent_in_sandbox(name: str) -> str:
        """Spin a microVM, ``pip install a2a-pack`` + the agent's own
        requirements.txt, then run ``a2a card`` to verify the scaffold
        produces a valid Card. Returns stdout/stderr.
        """
        try:
            prefix = _agent_prefix(name)
        except ValueError as exc:
            return json.dumps({"error": str(exc)})

        # Bundle the workspace files into the sandbox via base64 tar so
        # the sandbox doesn't need MinIO read access for this skill.
        bundle_bytes = _tarball_workspace_dir(s3, bucket, prefix)
        if not bundle_bytes:
            return json.dumps({
                "error": "no files in agents/" + name + "/ — write some first",
            })
        b64 = base64.b64encode(bundle_bytes).decode("ascii")
        script = (
            "set -e\n"
            "pip install --quiet a2a-pack >/dev/null\n"
            "mkdir -p /tmp/agent\n"
            f"echo '{b64}' | base64 -d | tar -xzf - -C /tmp/agent\n"
            "cd /tmp/agent\n"
            "ls -la\n"
            "if [ -f requirements.txt ]; then\n"
            "  pip install --quiet -r requirements.txt >/dev/null || true\n"
            "fi\n"
            "echo '--- card ---'\n"
            "a2a card --project .\n"
        )
        async with httpx.AsyncClient(timeout=settings.sandbox_timeout_s + 30) as c:
            r = await c.post(
                f"{settings.sandbox_url}/v1/run_shell",
                json={
                    "bucket": bucket, "script": script,
                    "image": settings.image, "memory_mib": 512,
                    "timeout_seconds": settings.sandbox_timeout_s,
                },
            )
        if r.status_code >= 400:
            return json.dumps({"error": f"sandbox {r.status_code}", "detail": r.text[:1000]})
        out = r.json()
        return json.dumps({
            "exit_code": out.get("exit_code"),
            "stdout": (out.get("stdout") or "")[:4000],
            "stderr": (out.get("stderr") or "")[:4000],
        })

    @tool
    async def cp_deploy_tarball(name: str, version: str = "0.1.0", public: bool = True) -> str:
        """Tar the agent at ``agents/<name>/`` and POST it to the control
        plane's ``/v1/agents/from-tarball`` endpoint on the user's behalf,
        then BLOCK until the live ``.well-known/agent-card`` reports the
        new version (or 180s timeout).

        Returns JSON with ``ok``, ``live``, ``url``, ``version``,
        ``live_version``, and ``live_skills[{name, input_schema}]``. When
        ``live=true`` the agent is callable; when ``live=false`` the
        deploy POST succeeded but the pod never came up with the new
        version — caller should treat this as a failure, not success.

        Compare ``live_skills[].input_schema`` against the code you wrote
        to confirm the advertised schema matches the params the function
        actually accepts. Mismatches cause downstream ``agent 400`` from
        ``call_agent``.

        Requires the orchestrator to have forwarded the user's CP JWT
        (the platform does this automatically when this agent's Card
        declared ``wants_cp_jwt=True``).
        """
        try:
            prefix = _agent_prefix(name)
        except ValueError as exc:
            return json.dumps({"error": str(exc)})
        if not ctx.cp_jwt:
            return json.dumps({
                "error": "no CP JWT forwarded — agent declaration missing wants_cp_jwt=True",
            })

        bundle_bytes = _tarball_workspace_dir(s3, bucket, prefix)
        if not bundle_bytes:
            return json.dumps({"error": f"no files at agents/{name}/"})
        # Read the entrypoint from a2a.yaml so we can pass it verbatim.
        entrypoint = _read_entrypoint(bundle_bytes) or f"agent:{_class_name(name)}"

        # NOTE: the CP from-tarball endpoint names the file field
        # ``source`` (see apps/control-plane/control_plane/routes/agents.py).
        files = {"source": (f"{name}.tar.gz", bundle_bytes, "application/gzip")}
        data = {
            "name": name, "version": version,
            "entrypoint": entrypoint, "public": "true" if public else "false",
            "description": "Built by agent-builder.",
        }
        try:
            async with httpx.AsyncClient(timeout=60.0) as c:
                r = await c.post(
                    f"{settings.cp_url}/v1/agents/from-tarball",
                    headers={"authorization": f"bearer {ctx.cp_jwt}"},
                    files=files, data=data,
                )
        except httpx.HTTPError as exc:
            return json.dumps({"error": f"cp unreachable: {exc}"})
        if r.status_code >= 400:
            return json.dumps({"error": f"cp {r.status_code}", "detail": r.text[:1000]})
        body = r.json()
        name_ = body.get("name")
        version_ = body.get("version")
        url_ = body.get("url")

        # CP returns immediately with status="building". The agent isn't
        # actually callable until the pod rolls + serves the new card.
        # Poll /.well-known/agent-card until the live card reports the
        # version we just deployed (or timeout). Surface the live skills
        # so the LLM can verify the input_schema actually matches the
        # code it wrote — catches the schema/code skew that causes
        # downstream 400s on call_agent.
        live, live_card = await _wait_for_live_card(url_, version_)
        out: dict[str, Any] = {
            "ok": live,
            "name": name_,
            "version": version_,
            "status": body.get("status"),
            "url": url_,
            "live": live,
        }
        if live:
            out["live_version"] = live_card.get("version")
            out["live_skills"] = [
                {
                    "name": s.get("name"),
                    "input_schema": s.get("input_schema") or {},
                }
                for s in (live_card.get("skills") or [])
            ]
        else:
            out["error"] = (
                "deploy did not become live within timeout — pod may still"
                " be rolling or build failed. Call cp_deploy_tarball again"
                " or inspect via discover_agent before declaring done."
            )
            if live_card:
                out["last_seen_version"] = live_card.get("version")
        return json.dumps(out)

    @tool
    async def cp_refresh_agent(name: str) -> str:
        """Force the control plane to re-fetch the agent's live
        ``/.well-known/agent-card`` and update its stored copy.

        Use this when the pod was redeployed out-of-band (CI bump, manual
        rollout) and the CP's cached card is stale — symptoms are stale
        ``input_schema`` in ``list_my_agents`` / ``discover_agent`` and
        4xx from ``call_agent`` on what should be a valid skill call.

        ``cp_deploy_tarball`` already does this automatically by polling
        the live card, so you only need ``cp_refresh_agent`` when you
        did NOT deploy through this tool.

        Returns JSON: the fresh ``AgentDetailOut`` (with ``card`` +
        ``skills`` + ``input_schema``), or ``{error, ...}``.
        """
        if not ctx.cp_jwt:
            return json.dumps({
                "error": "no CP JWT forwarded — agent declaration missing wants_cp_jwt=True",
            })
        try:
            async with httpx.AsyncClient(timeout=15.0) as c:
                r = await c.get(
                    f"{settings.cp_url}/v1/agents/{name}",
                    headers={"authorization": f"bearer {ctx.cp_jwt}"},
                )
        except httpx.HTTPError as exc:
            return json.dumps({"error": f"cp unreachable: {exc}"})
        if r.status_code >= 400:
            return json.dumps({"error": f"cp {r.status_code}", "detail": r.text[:1000]})
        body = r.json() or {}
        return json.dumps({
            "ok": True,
            "name": body.get("name"),
            "version": body.get("version"),
            "status": body.get("status"),
            "url": body.get("url"),
            "card": body.get("card") or {},
        })

    @tool
    def list_a2a_pack(subdir: str = "") -> str:
        """List ``.py`` files under the installed ``a2a_pack`` package.

        Use this when you need to recall what the SDK actually exposes
        — what decorators exist, what ``RunContext`` methods are
        available, how ``WorkspaceClient`` is shaped, what the
        ``AgentRuntime`` fields are, etc. The agent code you scaffold
        runs on this exact package, so reading it is the authoritative
        reference (the system-prompt examples are intentionally
        abridged).

        ``subdir`` is an optional path segment relative to the package
        root (``""``, ``"serve"``, ``"mcp"`` …). Returns JSON:
        ``{root, files: [{path, size}]}``.
        """
        try:
            import a2a_pack
            import os as _os
            base = _os.path.dirname(_os.path.abspath(a2a_pack.__file__))
        except Exception as exc:  # noqa: BLE001
            return json.dumps({"error": f"a2a_pack not importable: {exc}"})
        if subdir.startswith("/") or ".." in subdir.split("/"):
            return json.dumps({"error": "subdir must be relative; '..' disallowed"})
        root = _os.path.join(base, subdir) if subdir else base
        if not _os.path.commonpath([base, _os.path.abspath(root)]) == base:
            return json.dumps({"error": "subdir escapes package root"})
        if not _os.path.isdir(root):
            return json.dumps({"error": f"not a directory: {subdir}"})
        out: list[dict[str, Any]] = []
        for dirpath, _, filenames in _os.walk(root):
            for fname in filenames:
                if not fname.endswith(".py"):
                    continue
                full = _os.path.join(dirpath, fname)
                rel = _os.path.relpath(full, base)
                try:
                    size = _os.path.getsize(full)
                except OSError:
                    continue
                out.append({"path": rel, "size": size})
        out.sort(key=lambda r: r["path"])
        return json.dumps({"root": base, "files": out})

    @tool
    def read_a2a_pack(path: str) -> str:
        """Read a file from the installed ``a2a_pack`` package.

        Path is relative to the package root, e.g. ``agent.py``,
        ``context.py``, ``runtime.py``, ``workspace.py``, ``card.py``,
        ``mcp/server.py``. Returns JSON: ``{path, content}``. Capped
        at 64 KiB; reach for ``list_a2a_pack`` first if you need to
        chase a deeper file.
        """
        try:
            import a2a_pack
            import os as _os
            base = _os.path.dirname(_os.path.abspath(a2a_pack.__file__))
        except Exception as exc:  # noqa: BLE001
            return json.dumps({"error": f"a2a_pack not importable: {exc}"})
        if path.startswith("/") or ".." in path.split("/"):
            return json.dumps({"error": "path must be relative; '..' disallowed"})
        full = _os.path.abspath(_os.path.join(base, path))
        if _os.path.commonpath([base, full]) != base:
            return json.dumps({"error": "path escapes package root"})
        if not _os.path.isfile(full):
            return json.dumps({"error": f"not a file: {path}"})
        try:
            with open(full, "r", encoding="utf-8") as fh:
                content = fh.read(64 * 1024 + 1)
        except OSError as exc:
            return json.dumps({"error": str(exc)})
        truncated = len(content) > 64 * 1024
        if truncated:
            content = content[: 64 * 1024]
        return json.dumps({
            "path": path,
            "content": content,
            "truncated": truncated,
        })

    return [
        init_agent_template, list_agent_files, write_agent_file, read_agent_file,
        test_agent_in_sandbox, cp_deploy_tarball, cp_refresh_agent,
        list_a2a_pack, read_a2a_pack,
    ]


# --- helpers (not tools) ---


def _tarball_workspace_dir(s3: Any, bucket: str, prefix: str) -> bytes:
    """Pull every object under ``prefix`` from MinIO and produce a
    gzipped tarball whose entries are relative to ``prefix``."""
    paginator = s3.get_paginator("list_objects_v2")
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tf:
        for page in paginator.paginate(Bucket=bucket, Prefix=prefix):
            for obj in page.get("Contents") or ():
                key = obj["Key"]
                rel = key[len(prefix):]
                if not rel:
                    continue
                body = s3.get_object(Bucket=bucket, Key=key)["Body"].read()
                info = tarfile.TarInfo(name=rel)
                info.size = len(body)
                info.mode = 0o644
                tf.addfile(info, io.BytesIO(body))
    return buf.getvalue() if buf.getbuffer().nbytes > 0 else b""


def _read_entrypoint(bundle: bytes) -> str | None:
    try:
        with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf:
            for member in tf.getmembers():
                if member.name.endswith("a2a.yaml") and tf.extractfile(member):
                    text = tf.extractfile(member).read().decode("utf-8", "replace")  # type: ignore[union-attr]
                    for line in text.splitlines():
                        if line.lstrip().startswith("entrypoint:"):
                            return line.split(":", 1)[1].strip().strip('"\'')
    except Exception:  # noqa: BLE001
        pass
    return None


def _read_bundle_text(bundle: bytes, path: str) -> str | None:
    try:
        with tarfile.open(fileobj=io.BytesIO(bundle), mode="r:gz") as tf:
            for member in tf.getmembers():
                if member.name == path and member.isfile():
                    extracted = tf.extractfile(member)
                    if extracted is None:
                        return None
                    return extracted.read().decode("utf-8", "replace")
    except Exception:  # noqa: BLE001
        return None
    return None


def _class_name(slug: str) -> str:
    return "".join(p.capitalize() for p in re.split(r"[-_]+", slug) if p) or "MyAgent"


def _render_a2a_init_template(
    name: str,
    *,
    description: str = "A new A2A agent",
) -> dict[str, str]:
    _validate_name(name)
    class_name = _class_name(name)
    files = {
        "agent.py": _render_sdk_template(
            "agent.py.tmpl",
            name=name,
            class_name=class_name,
            description=description,
        ),
        "a2a.yaml": _render_sdk_template(
            "a2a.yaml.tmpl",
            name=name,
            class_name=class_name,
        ),
        "requirements.txt": _render_sdk_template("requirements.txt.tmpl"),
    }
    return files


def _render_sdk_template(template: str, /, **vars: str) -> str:
    text = (
        resources.files("a2a_pack.cli.templates")
        .joinpath(template)
        .read_text(encoding="utf-8")
    )
    for key, value in vars.items():
        text = text.replace("{{ " + key + " }}", value)
    return text