"""Persistent local sub-agent storage for the web UI."""

from __future__ import annotations

import json
import re
import shutil
from dataclasses import asdict, dataclass, field
from pathlib import Path
from typing import Any

from beaver.coordinator.registry import AgentRegistry


_INVALID_ID_RE = re.compile(r"[^a-z0-9-]+")


def normalize_subagent_id(value: str) -> str:
    normalized = _INVALID_ID_RE.sub("-", str(value or "").strip().lower()).strip("-")
    normalized = re.sub(r"-{2,}", "-", normalized)
    if not normalized:
        raise ValueError("Sub-agent id is required")
    return normalized


@dataclass(slots=True)
class SubagentSpec:
    id: str
    name: str
    description: str
    enabled: bool = True
    workspace: str = ""
    system_prompt: str = ""
    model: str | None = None
    delegation_mode: str = "remote_a2a_only"
    allow_mcp: bool = True
    tags: list[str] = field(default_factory=list)
    aliases: list[str] = field(default_factory=list)
    mcp_servers: dict[str, dict[str, Any]] = field(default_factory=dict)
    metadata: dict[str, Any] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, payload: dict[str, Any], *, workspace_path: Path | None = None) -> "SubagentSpec":
        agent_id = normalize_subagent_id(str(payload.get("id") or ""))
        name = str(payload.get("name") or agent_id).strip() or agent_id
        description = str(payload.get("description") or name).strip() or name
        workspace = str(payload.get("workspace") or "").strip()
        if not workspace and workspace_path is not None:
            workspace = str(workspace_path)
        mcp_servers = payload.get("mcp_servers", {})
        metadata = payload.get("metadata", {})
        return cls(
            id=agent_id,
            name=name,
            description=description,
            enabled=bool(payload.get("enabled", True)),
            workspace=workspace,
            system_prompt=str(payload.get("system_prompt") or "").strip(),
            model=(str(payload.get("model") or "").strip() or None),
            delegation_mode=str(payload.get("delegation_mode") or "remote_a2a_only").strip() or "remote_a2a_only",
            allow_mcp=bool(payload.get("allow_mcp", True)),
            tags=_string_list(payload.get("tags")),
            aliases=_string_list(payload.get("aliases")),
            mcp_servers=mcp_servers if isinstance(mcp_servers, dict) else {},
            metadata=metadata if isinstance(metadata, dict) else {},
        )

    def to_dict(self) -> dict[str, Any]:
        payload = asdict(self)
        if not self.model:
            payload["model"] = None
        return payload


class LocalSubagentStore:
    """Persist sub-agent definitions under `<workspace>/agents/<id>_agent/`."""

    def __init__(self, workspace: Path, *, public_base_url: str = "") -> None:
        self.workspace = workspace.expanduser().resolve()
        self.directory = self.workspace / "agents"
        self.public_base_url = public_base_url.rstrip("/")

    def list_subagents(self) -> list[SubagentSpec]:
        if not self.directory.exists():
            return []
        result: list[SubagentSpec] = []
        for child in sorted(self.directory.iterdir()):
            agents_json = child / "AGENTS.json"
            if not child.is_dir() or not agents_json.exists():
                continue
            try:
                payload = json.loads(agents_json.read_text(encoding="utf-8"))
            except (OSError, json.JSONDecodeError, ValueError):
                continue
            if isinstance(payload, dict):
                result.append(SubagentSpec.from_dict(payload, workspace_path=child))
        return result

    def get_subagent(self, agent_id: str) -> SubagentSpec | None:
        path = self.agents_json_path(agent_id)
        if not path.exists():
            return None
        try:
            payload = json.loads(path.read_text(encoding="utf-8"))
        except (OSError, json.JSONDecodeError, ValueError):
            return None
        if not isinstance(payload, dict):
            return None
        return SubagentSpec.from_dict(payload, workspace_path=self.subagent_dir(agent_id))

    def upsert_subagent(self, payload: dict[str, Any]) -> SubagentSpec:
        agent_id = normalize_subagent_id(str(payload.get("id") or ""))
        workspace_path = self.subagent_dir(agent_id)
        spec = SubagentSpec.from_dict(payload, workspace_path=workspace_path)
        self._ensure_workspace(workspace_path)
        spec.workspace = str(workspace_path)
        self._sync_agents_md(workspace_path, spec)
        self.agents_json_path(agent_id).write_text(
            json.dumps(spec.to_dict(), indent=2, ensure_ascii=False) + "\n",
            encoding="utf-8",
        )
        AgentRegistry(self.workspace).upsert_agent(self.build_registry_record(spec))
        return spec

    def delete_subagent(self, agent_id: str) -> bool:
        agent_id = normalize_subagent_id(agent_id)
        target = self.subagent_dir(agent_id)
        if not target.exists():
            return False
        AgentRegistry(self.workspace).delete_agent(agent_id)
        shutil.rmtree(target)
        return True

    def subagent_dir(self, agent_id: str) -> Path:
        return self.directory / f"{normalize_subagent_id(agent_id)}_agent"

    def agents_json_path(self, agent_id: str) -> Path:
        return self.subagent_dir(agent_id) / "AGENTS.json"

    def local_base_url(self, agent_id: str) -> str:
        if self.public_base_url:
            return f"{self.public_base_url}/subagents/{normalize_subagent_id(agent_id)}"
        return f"/subagents/{normalize_subagent_id(agent_id)}"

    def build_registry_record(self, spec: SubagentSpec) -> dict[str, Any]:
        base_url = self.local_base_url(spec.id)
        return {
            "agent_id": spec.id,
            "name": spec.id,
            "display_name": spec.name,
            "role": spec.description,
            "description": spec.description,
            "system_prompt": spec.system_prompt,
            "model": spec.model,
            "tags": sorted(set(["local-subagent", *spec.tags])),
            "status": "active" if spec.enabled else "disabled",
            "source": "workspace",
            "metadata": {
                **spec.metadata,
                "workspace": spec.workspace,
                "managed_by": "subagent-manager",
                "local_subagent": True,
                "kind": "local_subagent",
                "protocol": "a2a",
                "base_url": base_url,
                "endpoint": f"{base_url}/rpc",
                "card_url": f"{base_url}/.well-known/agent-card",
                "aliases": sorted(set([spec.name, *spec.aliases])),
            },
        }

    def serialize(self, spec: SubagentSpec) -> dict[str, Any]:
        base_url = self.local_base_url(spec.id)
        return {
            **spec.to_dict(),
            "base_url": base_url,
            "endpoint": f"{base_url}/rpc",
            "card_url": f"{base_url}/.well-known/agent-card",
        }

    def _ensure_workspace(self, workspace_path: Path) -> None:
        workspace_path.mkdir(parents=True, exist_ok=True)
        (workspace_path / "memory").mkdir(exist_ok=True)
        (workspace_path / "skills").mkdir(exist_ok=True)

    def _sync_agents_md(self, workspace_path: Path, spec: SubagentSpec) -> None:
        (workspace_path / "AGENTS.md").write_text(self._render_agents_md(spec), encoding="utf-8")

    @staticmethod
    def _render_agents_md(spec: SubagentSpec) -> str:
        prompt = spec.system_prompt.strip() or "Complete delegated tasks accurately and concisely."
        return f"""# {spec.name}

You are {spec.name}, a persistent local sub-agent managed by Beaver.

## Role
{spec.description}

## System Prompt
{prompt}

## Constraints
- Work only inside this workspace.
- Respond only to delegated tasks.
- Do not create or manage local sub-agents.
- Do not message end users directly.
"""


def _string_list(value: Any) -> list[str]:
    if isinstance(value, str):
        value = [item.strip() for item in value.split(",")]
    if not isinstance(value, list):
        return []
    result: list[str] = []
    for item in value:
        text = str(item).strip()
        if text and text not in result:
            result.append(text)
    return result
