"""
Thin client for talking to the OffSec Shield Node (offsec-agents MCP backend).

Usage:

    from scripts.offsec_node_client import OffsecNodeClient

    client = OffsecNodeClient()  # uses OFFSEC_NODE_URL env
    agents = await client.command("agents list")
    status = await client.command("tem status")

"""

from __future__ import annotations

import asyncio
import json
import os
from dataclasses import dataclass
from typing import Any, Dict, Optional

import aiohttp


DEFAULT_OFFSEC_NODE_URL = "http://shield-vm:8081"


@dataclass
class OffsecNodeError(Exception):
    message: str
    status: Optional[int] = None
    details: Optional[Dict[str, Any]] = None

    def __str__(self) -> str:
        base = self.message
        if self.status is not None:
            base += f" (status={self.status})"
        if self.details:
            base += f" details={self.details}"
        return base


@dataclass
class OffsecNodeClient:
    base_url: str = DEFAULT_OFFSEC_NODE_URL
    timeout_seconds: int = 10

    @classmethod
    def from_env(cls) -> "OffsecNodeClient":
        url = os.getenv("OFFSEC_NODE_URL", DEFAULT_OFFSEC_NODE_URL)
        return cls(base_url=url)

    async def health(self) -> Dict[str, Any]:
        url = f"{self.base_url.rstrip('/')}/health"
        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout_seconds)) as session:
            async with session.get(url) as resp:
                if resp.status != 200:
                    raise OffsecNodeError("Shield node health check failed", status=resp.status)
                return await resp.json()

    async def command(
        self,
        command: str,
        session_id: str = "vaultmesh-client",
        user: str = "vaultmesh",
    ) -> Dict[str, Any]:
        """
        Send a command to the Shield Node MCP backend.

        Example commands:
            "agents list"
            "status"
            "shield status"
            "proof latest"
            "agent spawn recon"
        """
        url = f"{self.base_url.rstrip('/')}/mcp/command"
        payload: Dict[str, Any] = {
            "session_id": session_id,
            "user": user,
            "command": command,
        }

        async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout_seconds)) as session:
            async with session.post(url, json=payload) as resp:
                text = await resp.text()
                if resp.status != 200:
                    # Try to parse JSON details if present
                    try:
                        data = json.loads(text)
                    except json.JSONDecodeError:
                        data = None
                    raise OffsecNodeError(
                        "Shield node command failed",
                        status=resp.status,
                        details=data or {"raw": text},
                    )
                try:
                    return json.loads(text)
                except json.JSONDecodeError as exc:
                    raise OffsecNodeError("Invalid JSON from Shield node", details={"raw": text}) from exc


# Optional: CLI entrypoint for quick manual tests
async def _demo() -> None:
    client = OffsecNodeClient.from_env()
    print(f"[offsec-node] base_url={client.base_url}")
    try:
        health = await client.health()
        print("Health:", json.dumps(health, indent=2))
    except OffsecNodeError as e:
        print("Health check failed:", e)
        return

    try:
        agents = await client.command("agents list")
        print("Agents:", json.dumps(agents, indent=2))
    except OffsecNodeError as e:
        print("Command failed:", e)


if __name__ == "__main__":
    asyncio.run(_demo())