"""Persistent autonomous tool synthesis for agents and teams."""
from __future__ import annotations

import importlib.util
import inspect
import json
import os
from dataclasses import dataclass
from datetime import UTC, datetime
from pathlib import Path
import re
import shutil
import subprocess
import sys
from typing import Any
from uuid import uuid4

from logos_adk.tools import Tool, tool
from logos_adk.tools.registry import tool_registry


_TOOL_NAME_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")
_DEFAULT_GENERATED_TOOLS_DIR = ".logos_adk/generated_tools"


@dataclass(frozen=True)
class GeneratedToolArtifact:
    """Persisted generated tool metadata."""

    name: str
    description: str
    parameters: dict[str, Any]
    module_path: str
    tests_path: str
    created_at: str
    created_by: str | None = None


class GeneratedToolStore:
    """Filesystem-backed store for synthesized tools and their bindings."""

    def __init__(self, root: str | os.PathLike[str] | None = None) -> None:
        configured_root = root or os.getenv("LOGOS_ADK_GENERATED_TOOLS_DIR") or _DEFAULT_GENERATED_TOOLS_DIR
        root_path = Path(configured_root)
        if not root_path.is_absolute():
            root_path = Path.cwd() / root_path
        self.root = root_path

    def save_tool(
        self,
        *,
        name: str,
        description: str,
        parameters: dict[str, Any],
        source_code: str,
        tests: str,
        created_by: str | None = None,
    ) -> GeneratedToolArtifact:
        """Persist a generated tool after validating its code and tests."""
        tool_name = _normalize_tool_name(name)
        tool_description = _normalize_description(description)
        parameter_schema = _normalize_parameter_schema(parameters)
        if not isinstance(source_code, str) or not source_code.strip():
            raise ValueError("source_code must be a non-empty string")
        if not isinstance(tests, str) or not tests.strip():
            raise ValueError("tests must be a non-empty string")

        created_at = datetime.now(UTC).isoformat()
        artifact = GeneratedToolArtifact(
            name=tool_name,
            description=tool_description,
            parameters=parameter_schema,
            module_path=str(self._tool_dir(tool_name) / "tool.py"),
            tests_path=str(self._tool_dir(tool_name) / "tests.py"),
            created_at=created_at,
            created_by=created_by,
        )
        manifest_payload = {
            "name": artifact.name,
            "description": artifact.description,
            "parameters": artifact.parameters,
            "module_path": "tool.py",
            "tests_path": "tests.py",
            "created_at": artifact.created_at,
            "created_by": artifact.created_by,
        }

        self.root.mkdir(parents=True, exist_ok=True)
        staging_dir = self.root / f".staging-{tool_name}-{uuid4().hex}"
        staging_dir.mkdir(parents=True, exist_ok=False)
        try:
            module_path = staging_dir / "tool.py"
            tests_path = staging_dir / "tests.py"
            manifest_path = staging_dir / "manifest.json"
            module_path.write_text(source_code, encoding="utf-8")
            tests_path.write_text(tests, encoding="utf-8")
            manifest_path.write_text(
                json.dumps(manifest_payload, indent=2, ensure_ascii=False, sort_keys=True),
                encoding="utf-8",
            )
            self._validate_candidate(module_path, tests_path)
            final_dir = self._tool_dir(tool_name)
            if final_dir.exists():
                shutil.rmtree(final_dir)
            staging_dir.rename(final_dir)
        finally:
            if staging_dir.exists():
                shutil.rmtree(staging_dir, ignore_errors=True)
        return self.get_artifact(tool_name)

    def get_artifact(self, name: str) -> GeneratedToolArtifact:
        """Return stored metadata for a generated tool."""
        tool_name = _normalize_tool_name(name)
        manifest_path = self._tool_dir(tool_name) / "manifest.json"
        if not manifest_path.exists():
            raise FileNotFoundError(f"Generated tool '{tool_name}' not found in store '{self.root}'")
        payload = json.loads(manifest_path.read_text(encoding="utf-8"))
        return GeneratedToolArtifact(
            name=payload["name"],
            description=payload["description"],
            parameters=payload["parameters"],
            module_path=str(self._tool_dir(tool_name) / payload.get("module_path", "tool.py")),
            tests_path=str(self._tool_dir(tool_name) / payload.get("tests_path", "tests.py")),
            created_at=payload["created_at"],
            created_by=payload.get("created_by"),
        )

    def list_tools(self) -> list[str]:
        """List stored generated tool names."""
        if not self.root.exists():
            return []
        names: list[str] = []
        for entry in self.root.iterdir():
            if not entry.is_dir() or entry.name.startswith("."):
                continue
            if (entry / "manifest.json").exists():
                names.append(entry.name)
        return sorted(names)

    def resolve_tool(self, name: str) -> Tool | None:
        """Load and materialize a persisted generated tool."""
        try:
            artifact = self.get_artifact(name)
        except FileNotFoundError:
            return None
        tool_obj = self._build_tool(artifact)
        tool_registry.register(tool_obj.name, tool_obj)
        return tool_obj

    def bind_tool_to_agent(self, tool_name: str, agent_name: str) -> None:
        """Persist an agent-level binding for a generated tool."""
        binding_name = _normalize_binding_name(agent_name, kind="agent")
        normalized_tool_name = _normalize_tool_name(tool_name)
        bindings = self._load_bindings()
        agent_tools = bindings.setdefault("agents", {})
        tool_names = agent_tools.setdefault(binding_name, [])
        if normalized_tool_name not in tool_names:
            tool_names.append(normalized_tool_name)
        self._write_bindings(bindings)

    def bind_tool_to_team(self, tool_name: str, team_name: str) -> None:
        """Persist a team-level binding for a generated tool."""
        binding_name = _normalize_binding_name(team_name, kind="team")
        normalized_tool_name = _normalize_tool_name(tool_name)
        bindings = self._load_bindings()
        team_tools = bindings.setdefault("teams", {})
        tool_names = team_tools.setdefault(binding_name, [])
        if normalized_tool_name not in tool_names:
            tool_names.append(normalized_tool_name)
        self._write_bindings(bindings)

    def attach_tools_for_agent(self, agent: Any) -> list[Tool]:
        """Attach all persisted agent-bound generated tools."""
        agent_name = _normalize_binding_name(getattr(agent, "name", None), kind="agent")
        bindings = self._load_bindings()
        attached: list[Tool] = []
        for tool_name in bindings.get("agents", {}).get(agent_name, []):
            tool_obj = self.resolve_tool(tool_name)
            if tool_obj is None:
                continue
            agent.tools([tool_obj])
            attached.append(tool_obj)
        return attached

    def attach_tools_for_team(self, team: Any) -> list[Tool]:
        """Attach all persisted team-bound generated tools to every team member."""
        team_name = _normalize_binding_name(getattr(team, "name", None), kind="team")
        bindings = self._load_bindings()
        attached: list[Tool] = []
        for tool_name in bindings.get("teams", {}).get(team_name, []):
            tool_obj = self.resolve_tool(tool_name)
            if tool_obj is None:
                continue
            for member in getattr(team, "_agents", {}).values():
                member.tools([tool_obj])
            attached.append(tool_obj)
        return attached

    def _tool_dir(self, name: str) -> Path:
        return self.root / _normalize_tool_name(name)

    def _bindings_path(self) -> Path:
        return self.root / "bindings.json"

    def _load_bindings(self) -> dict[str, dict[str, list[str]]]:
        path = self._bindings_path()
        if not path.exists():
            return {"agents": {}, "teams": {}}
        payload = json.loads(path.read_text(encoding="utf-8"))
        if not isinstance(payload, dict):
            return {"agents": {}, "teams": {}}
        agents = payload.get("agents")
        teams = payload.get("teams")
        return {
            "agents": agents if isinstance(agents, dict) else {},
            "teams": teams if isinstance(teams, dict) else {},
        }

    def _write_bindings(self, bindings: dict[str, dict[str, list[str]]]) -> None:
        self.root.mkdir(parents=True, exist_ok=True)
        self._bindings_path().write_text(
            json.dumps(bindings, indent=2, ensure_ascii=False, sort_keys=True),
            encoding="utf-8",
        )

    def _validate_candidate(self, module_path: Path, tests_path: Path) -> None:
        self._load_module_from_path(module_path, module_name=f"logos_adk_generated_probe_{module_path.stem}")
        project_root = Path(__file__).resolve().parent.parent
        env = dict(os.environ)
        python_path_entries = [str(project_root)]
        if env.get("PYTHONPATH"):
            python_path_entries.append(env["PYTHONPATH"])
        env["PYTHONPATH"] = os.pathsep.join(python_path_entries)
        harness = (
            "from __future__ import annotations\n"
            "import asyncio\n"
            "import importlib.util\n"
            "import inspect\n"
            "import pathlib\n"
            "import sys\n"
            "module_path = pathlib.Path(sys.argv[1])\n"
            "tests_path = pathlib.Path(sys.argv[2])\n"
            "spec = importlib.util.spec_from_file_location('generated_tool_module', module_path)\n"
            "if spec is None or spec.loader is None:\n"
            "    raise RuntimeError(f'Unable to import generated tool module from {module_path}')\n"
            "module = importlib.util.module_from_spec(spec)\n"
            "spec.loader.exec_module(module)\n"
            "namespace = {'module': module, '__name__': '__main__'}\n"
            "code = compile(tests_path.read_text(encoding='utf-8'), str(tests_path), 'exec')\n"
            "exec(code, namespace, namespace)\n"
            "candidate = namespace.get('run_tests') or namespace.get('test')\n"
            "if candidate is not None:\n"
            "    result = candidate()\n"
            "    if inspect.isawaitable(result):\n"
            "        asyncio.run(result)\n"
        )
        completed = subprocess.run(
            [sys.executable, "-c", harness, str(module_path), str(tests_path)],
            cwd=str(module_path.parent),
            env=env,
            capture_output=True,
            text=True,
            check=False,
        )
        if completed.returncode != 0:
            details = completed.stderr.strip() or completed.stdout.strip() or "generated tool tests failed"
            raise ValueError(f"Generated tool validation failed: {details}")

    def _build_tool(self, artifact: GeneratedToolArtifact) -> Tool:
        module = self._load_module_from_path(Path(artifact.module_path), module_name=f"logos_adk_generated_{artifact.name}")
        run_fn = getattr(module, "run", None)
        if run_fn is None or not callable(run_fn):
            raise ValueError(
                f"Generated tool '{artifact.name}' must define callable 'run(**kwargs)' in {artifact.module_path}"
            )

        async def invoke(**kwargs: Any) -> Any:
            result = run_fn(**kwargs)
            if inspect.isawaitable(result):
                result = await result
            return result

        return Tool(
            fn=invoke,
            name=artifact.name,
            description=artifact.description,
            schema={
                "name": artifact.name,
                "description": artifact.description,
                "parameters": artifact.parameters,
            },
        )

    def _load_module_from_path(self, module_path: Path, *, module_name: str) -> Any:
        if not module_path.exists():
            raise FileNotFoundError(f"Generated tool module not found: {module_path}")
        spec = importlib.util.spec_from_file_location(f"{module_name}_{uuid4().hex}", module_path)
        if spec is None or spec.loader is None:
            raise ValueError(f"Unable to load module from {module_path}")
        module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(module)
        return module


generated_tool_store = GeneratedToolStore()


def _normalize_tool_name(name: str) -> str:
    if not isinstance(name, str) or not name.strip():
        raise ValueError("Tool name must be a non-empty string")
    normalized = name.strip()
    if not _TOOL_NAME_RE.match(normalized):
        raise ValueError(
            "Tool name must be a valid Python identifier containing only letters, numbers, and underscores"
        )
    return normalized


def _normalize_description(description: str) -> str:
    if not isinstance(description, str) or not description.strip():
        raise ValueError("Tool description must be a non-empty string")
    return description.strip()


def _normalize_parameter_schema(parameters: dict[str, Any]) -> dict[str, Any]:
    if not isinstance(parameters, dict):
        raise ValueError("parameters must be a JSON-schema object")
    schema_type = parameters.get("type")
    properties = parameters.get("properties")
    if schema_type != "object" or not isinstance(properties, dict):
        raise ValueError("parameters must look like a JSON schema object with type='object' and properties")
    required = parameters.get("required", [])
    if not isinstance(required, list) or not all(isinstance(item, str) for item in required):
        raise ValueError("parameters.required must be a list of strings")
    return json.loads(json.dumps(parameters))


def _normalize_binding_name(name: Any, *, kind: str) -> str:
    if not isinstance(name, str) or not name.strip():
        raise ValueError(f"{kind} name must be a non-empty string")
    return name.strip()


def _resolve_creator_agent(candidate: Any = None) -> Any | None:
    if candidate is not None and hasattr(candidate, "_config") and hasattr(candidate, "name"):
        return candidate
    frame = inspect.currentframe()
    while frame is not None:
        obj = frame.f_locals.get("self")
        if obj is not None and hasattr(obj, "_config") and hasattr(obj, "name"):
            return obj
        frame = frame.f_back
    return None


@tool(
    description=(
        "Develop a new Python tool from source code and tests, validate it in a sandboxed subprocess, "
        "persist it on disk, and attach it to the current agent or team."
    )
)
async def develop_tool(
    name: str,
    description: str,
    parameters: dict[str, Any],
    source_code: str,
    tests: str,
    scope: str = "agent",
    store_path: str | None = None,
    self: Any = None,
) -> dict[str, Any]:
    """Create, validate, persist, and install a new tool.

    The generated Python module must define callable ``run(**kwargs)``.
    The tests body can use ``module`` and may optionally define ``run_tests()`` or ``test()``.
    """
    creator = _resolve_creator_agent(self)
    normalized_scope = scope.strip().lower() if isinstance(scope, str) else ""
    if normalized_scope not in {"agent", "team"}:
        raise ValueError("scope must be either 'agent' or 'team'")
    if creator is None:
        raise RuntimeError("develop_tool must be called from an Agent context or with self=<agent>")
    if normalized_scope == "team" and getattr(creator, "_team", None) is None:
        raise RuntimeError("scope='team' requires the creator agent to belong to a team")

    store = GeneratedToolStore(store_path)
    artifact = store.save_tool(
        name=name,
        description=description,
        parameters=parameters,
        source_code=source_code,
        tests=tests,
        created_by=getattr(creator, "name", None),
    )
    tool_obj = store.resolve_tool(artifact.name)
    if tool_obj is None:
        raise RuntimeError(f"Generated tool '{artifact.name}' could not be loaded after persistence")

    installed_on: list[str] = []
    if normalized_scope == "agent":
        store.bind_tool_to_agent(artifact.name, creator.name)
        creator.tools([tool_obj])
        installed_on.append(creator.name)
    else:
        team = creator._team
        store.bind_tool_to_team(artifact.name, team.name)
        for member in team._agents.values():
            member.tools([tool_obj])
            installed_on.append(member.name)

    return {
        "status": "installed",
        "tool_name": artifact.name,
        "scope": normalized_scope,
        "installed_on": installed_on,
        "description": artifact.description,
        "module_path": artifact.module_path,
        "tests_path": artifact.tests_path,
        "created_at": artifact.created_at,
    }


__all__ = [
    "GeneratedToolArtifact",
    "GeneratedToolStore",
    "develop_tool",
    "generated_tool_store",
]
