import httpx
from typing import Optional
from app.config import settings
from app.internal.logger import logger


class C2Client:
    def __init__(self):
        self.base = settings.c2_url.rstrip("/")

    def _headers(self) -> dict:
        if settings.c2_service_key:
            return {"Authorization": f"Bearer {settings.c2_service_key}"}
        return {}

    async def get_nodes(self) -> list:
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                r = await client.get(f"{self.base}/nodes", headers=self._headers())
                r.raise_for_status()
                return r.json()
        except Exception as e:
            logger.error(f"C2 get_nodes failed: {e}")
            return []

    async def get_systems(self) -> list:
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                r = await client.get(f"{self.base}/systems", headers=self._headers())
                r.raise_for_status()
                return r.json()
        except Exception as e:
            logger.error(f"C2 get_systems failed: {e}")
            return []

    async def get_tokens(self) -> list:
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                r = await client.get(f"{self.base}/tokens", headers=self._headers())
                r.raise_for_status()
                return r.json()
        except Exception as e:
            logger.error(f"C2 get_tokens failed: {e}")
            return []

    async def send_command(self, node_id: str, payload: dict) -> bool:
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                r = await client.post(
                    f"{self.base}/nodes/{node_id}/command",
                    json=payload,
                    headers=self._headers(),
                )
                r.raise_for_status()
                return True
        except Exception as e:
            logger.error(f"C2 send_command failed: {e}")
            return False

    async def find_node_for_system(self, system_id: str) -> Optional[dict]:
        """Return the first online node assigned to the given system."""
        nodes = await self.get_nodes()
        for node in nodes:
            if (
                node.get("assigned_system_id") == system_id
                and node.get("status") in ("online", "recording")
            ):
                return node
        return None


c2 = C2Client()