#!/usr/bin/env python3
"""Manage persistent local sub-agents."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
import sys
from typing import Any

ROOT = Path(__file__).resolve().parents[4]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from nanobot.agent.subagents import LocalSubagentStore, SubagentSpec
from nanobot.config.loader import load_config


def _store():
    config = load_config()
    return config, LocalSubagentStore(config.workspace_path)


def _print_json(payload: Any) -> None:
    print(json.dumps(payload, indent=2, ensure_ascii=False))


def _load_spec_or_die(store: LocalSubagentStore, agent_id: str) -> SubagentSpec:
    spec = store.get_subagent(agent_id)
    if spec is None:
        raise SystemExit(f"Sub-agent not found: {agent_id}")
    return spec


def _parse_key_values(items: list[str]) -> dict[str, str]:
    result: dict[str, str] = {}
    for item in items:
        if "=" not in item:
            raise SystemExit(f"Expected KEY=VALUE, got: {item}")
        key, value = item.split("=", 1)
        key = key.strip()
        if not key:
            raise SystemExit(f"Invalid empty key in: {item}")
        result[key] = value
    return result


def cmd_list(_: argparse.Namespace) -> None:
    _, store = _store()
    _print_json([spec.to_dict() for spec in store.list_subagents()])


def cmd_show(args: argparse.Namespace) -> None:
    _, store = _store()
    spec = _load_spec_or_die(store, args.agent_id)
    _print_json(spec.to_dict())


def cmd_create(args: argparse.Namespace) -> None:
    config, store = _store()
    current = store.get_subagent(args.agent_id)
    payload = current.to_dict() if current is not None else {"id": args.agent_id}
    payload.update({
        "id": args.agent_id,
        "name": args.name or payload.get("name") or args.agent_id,
        "description": args.description or payload.get("description") or args.name or args.agent_id,
        "enabled": not args.disabled,
        "delegation_mode": payload.get("delegation_mode") or "remote_a2a_only",
    })
    if args.system_prompt:
        payload["system_prompt"] = args.system_prompt
    if args.model:
        payload["model"] = args.model
    spec = store.upsert_subagent(payload, config)
    _print_json(spec.to_dict())


def cmd_delete(args: argparse.Namespace) -> None:
    _, store = _store()
    deleted = store.delete_subagent(args.agent_id)
    if not deleted:
        raise SystemExit(f"Sub-agent not found: {args.agent_id}")
    _print_json({"ok": True, "id": args.agent_id})


def cmd_set_system_prompt(args: argparse.Namespace) -> None:
    config, store = _store()
    spec = _load_spec_or_die(store, args.agent_id)
    payload = spec.to_dict()
    payload["system_prompt"] = args.text
    updated = store.upsert_subagent(payload, config)
    _print_json(updated.to_dict())


def cmd_add_mcp_http(args: argparse.Namespace) -> None:
    config, store = _store()
    spec = _load_spec_or_die(store, args.agent_id)
    payload = spec.to_dict()
    payload.setdefault("mcp_servers", {})
    payload["mcp_servers"][args.server_id] = {
        "url": args.url,
        "headers": _parse_key_values(args.header),
        "auth_mode": args.auth_mode,
        "auth_audience": args.auth_audience,
        "auth_scopes": list(args.auth_scope),
        "tool_timeout": args.tool_timeout,
        "sensitive": args.sensitive,
    }
    updated = store.upsert_subagent(payload, config)
    _print_json(updated.to_dict())


def cmd_add_mcp_stdio(args: argparse.Namespace) -> None:
    config, store = _store()
    spec = _load_spec_or_die(store, args.agent_id)
    payload = spec.to_dict()
    payload.setdefault("mcp_servers", {})
    payload["mcp_servers"][args.server_id] = {
        "command": args.command,
        "args": list(args.arg),
        "env": _parse_key_values(args.env),
        "auth_mode": args.auth_mode,
        "auth_audience": args.auth_audience,
        "auth_scopes": list(args.auth_scope),
        "tool_timeout": args.tool_timeout,
        "sensitive": args.sensitive,
    }
    updated = store.upsert_subagent(payload, config)
    _print_json(updated.to_dict())


def cmd_remove_mcp(args: argparse.Namespace) -> None:
    config, store = _store()
    spec = _load_spec_or_die(store, args.agent_id)
    payload = spec.to_dict()
    mcp_servers = payload.setdefault("mcp_servers", {})
    mcp_servers.pop(args.server_id, None)
    updated = store.upsert_subagent(payload, config)
    _print_json(updated.to_dict())


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Manage persistent local sub-agents")
    sub = parser.add_subparsers(dest="command", required=True)

    list_parser = sub.add_parser("list", help="List sub-agents")
    list_parser.set_defaults(func=cmd_list)

    show_parser = sub.add_parser("show", help="Show one sub-agent")
    show_parser.add_argument("agent_id")
    show_parser.set_defaults(func=cmd_show)

    create_parser = sub.add_parser("create", help="Create or update a sub-agent")
    create_parser.add_argument("--id", dest="agent_id", required=True)
    create_parser.add_argument("--name", default="")
    create_parser.add_argument("--description", default="")
    create_parser.add_argument("--system-prompt", default="")
    create_parser.add_argument("--model", default="")
    create_parser.add_argument("--disabled", action="store_true")
    create_parser.set_defaults(func=cmd_create)

    delete_parser = sub.add_parser("delete", help="Delete a sub-agent")
    delete_parser.add_argument("agent_id")
    delete_parser.set_defaults(func=cmd_delete)

    prompt_parser = sub.add_parser("set-system-prompt", help="Update the system prompt")
    prompt_parser.add_argument("agent_id")
    prompt_parser.add_argument("--text", required=True)
    prompt_parser.set_defaults(func=cmd_set_system_prompt)

    http_parser = sub.add_parser("add-mcp-http", help="Add an HTTP MCP server")
    http_parser.add_argument("agent_id")
    http_parser.add_argument("--server-id", required=True)
    http_parser.add_argument("--url", required=True)
    http_parser.add_argument("--header", action="append", default=[])
    http_parser.add_argument("--auth-mode", default="none")
    http_parser.add_argument("--auth-audience", default="")
    http_parser.add_argument("--auth-scope", action="append", default=[])
    http_parser.add_argument("--tool-timeout", type=int, default=30)
    http_parser.add_argument("--sensitive", action="store_true")
    http_parser.set_defaults(func=cmd_add_mcp_http)

    stdio_parser = sub.add_parser("add-mcp-stdio", help="Add a stdio MCP server")
    stdio_parser.add_argument("agent_id")
    stdio_parser.add_argument("--server-id", required=True)
    stdio_parser.add_argument("--command", required=True)
    stdio_parser.add_argument("--arg", action="append", default=[])
    stdio_parser.add_argument("--env", action="append", default=[])
    stdio_parser.add_argument("--auth-mode", default="none")
    stdio_parser.add_argument("--auth-audience", default="")
    stdio_parser.add_argument("--auth-scope", action="append", default=[])
    stdio_parser.add_argument("--tool-timeout", type=int, default=30)
    stdio_parser.add_argument("--sensitive", action="store_true")
    stdio_parser.set_defaults(func=cmd_add_mcp_stdio)

    remove_mcp = sub.add_parser("remove-mcp", help="Remove an MCP server")
    remove_mcp.add_argument("agent_id")
    remove_mcp.add_argument("--server-id", required=True)
    remove_mcp.set_defaults(func=cmd_remove_mcp)

    return parser


def main() -> None:
    parser = build_parser()
    args = parser.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()