from __future__ import annotations

import io
import tarfile
import unittest

from agent_builder.tools import (
    _render_a2a_init_template,
)


class TemplateInitTests(unittest.TestCase):
    def test_render_a2a_init_template_matches_sdk_starter(self) -> None:
        files = _render_a2a_init_template(
            "research-agent",
            description="Research helper",
        )

        self.assertEqual(set(files), {"agent.py", "a2a.yaml", "requirements.txt"})
        self.assertIn('name = "research-agent"', files["agent.py"])
        self.assertIn('description = "Research helper"', files["agent.py"])
        self.assertIn("LLMProvisioning.CALLER_PROVIDED", files["agent.py"])
        self.assertIn("ctx.llm", files["agent.py"])
        self.assertIn("ctx.workspace_backend()", files["agent.py"])
        self.assertIn("create_deep_agent", files["agent.py"])
        self.assertIn("name: research-agent", files["a2a.yaml"])
        self.assertIn("entrypoint: agent:ResearchAgent", files["a2a.yaml"])
        self.assertIn("deepagents>=0.5.0", files["requirements.txt"])

    def test_render_a2a_init_template_rejects_invalid_names(self) -> None:
        with self.assertRaises(ValueError):
            _render_a2a_init_template("Bad Name")

    def test_agent_projects_can_diverge_from_template(self) -> None:
        files = _render_a2a_init_template("regex-explainer")
        files["agent.py"] = "from a2a_pack import A2AAgent\n"

        self.assertIn("A2AAgent", files["agent.py"])


def _tarball(files: dict[str, str]) -> bytes:
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tf:
        for path, text in files.items():
            body = text.encode("utf-8")
            info = tarfile.TarInfo(path)
            info.size = len(body)
            tf.addfile(info, io.BytesIO(body))
    return buf.getvalue()