import json
import logging
from typing import Dict, List, Optional

import httpx
import pandas as pd

logger = logging.getLogger(__name__)

class AirlockAPIWrapper:
    def __init__(self, base_url: str, api_key: str):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.headers = {"X-APIKey": self.api_key}

    async def _post(self, endpoint: str, payload: Optional[dict] = None) -> dict:
        url = f"{self.base_url}{endpoint}"
        data = json.dumps(payload or {})
        timeout = httpx.Timeout(300.0)
        async with httpx.AsyncClient(verify=False, timeout=timeout) as client:
            try:
                logger.debug(f"POST Request to {url} with payload: {payload}")
                response = await client.post(url, headers=self.headers, data=data) # pyright: ignore[reportArgumentType]
                response.raise_for_status()
                logger.debug(f"Response received from {url}")
                return response.json()
            except httpx.RequestError as e:
                logger.error(f"API request failed: {e}")
                raise

    # Allowlist Management
    async def allowlist_find_all(self) -> pd.DataFrame:
        result = await self._post("/v1/application", {})
        return pd.DataFrame(result["response"]["applications"])

    # Agent Management
    async def agent_find_all(self) -> pd.DataFrame:
        result = await self._post("/v1/agent/find", {})
        return pd.DataFrame(result["response"]["agents"])

    async def agent_find_by_hostname(self, hostname: str) -> pd.DataFrame:
        payload = {"hostname": hostname}
        result = await self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    async def agent_find_by_id(self, agentid: str) -> pd.DataFrame:
        payload = {"agentid": agentid}
        result = await self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    async def agent_find_by_status(self, status: int) -> pd.DataFrame:
        payload = {"status": status}
        result = await self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    async def agent_find_by_username(self, username: str) -> pd.DataFrame:
        payload = {"username": username}
        result = await self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    async def agent_move(self, agentid: str, groupid: str) -> dict:
        payload = {"agentid": agentid, "groupid": groupid}
        return await self._post("/v1/agent/move", payload)

    async def agents_find_by_group(self, groupid: str) -> pd.DataFrame:
        payload = {"groupid": groupid}
        result = await self._post("/v1/agent/find", payload)
        return pd.DataFrame(result["response"]["agents"])

    # Hash Management
    async def hash_add_to_allowlist(self, applicationid: str, hashes: List[str]) -> dict:
        payload = {"applicationid": applicationid, "hashes": hashes}
        return await self._post("/v1/hash/application/add", payload)

    async def hash_query(self, hashes: List[str]) -> pd.DataFrame:
        payload = {"hashes": hashes}
        result = await self._post("/v1/hash/query", payload)
        return pd.DataFrame(result["response"]["results"])

    # OTP Management
    async def otp_find_active(self) -> pd.DataFrame:
        payload = {"status": "1"}
        result = await self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    async def otp_find_awaiting(self) -> pd.DataFrame:
        payload = {"status": "0"}
        result = await self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    async def otp_find_by_agent(self, agentid) -> pd.DataFrame:
        payload = {"agentid": agentid}
        result = await self._post("/v1/otp/usage", payload)
        return pd.DataFrame(result["response"]["otpusage"])

    async def otp_generate(self, agentid: str, duration: int, purpose: str) -> str:
        payload = {
            "duration": str(duration),
            "agentid": str(agentid),
            "purpose": purpose,
        }
        result = await self._post("/v1/otp/retrieve", payload)
        return result["response"]["otpcode"]

    async def otp_get_activities(self, otpid: str) -> pd.DataFrame:
        payload = {"otpid": otpid}
        result = await self._post("/v1/otp/activities", payload)
        return pd.DataFrame(result["response"]["otpactivities"])

    async def otp_revoke(self, otpid: str) -> dict:
        payload = {"otpid": otpid}
        return await self._post("/v1/otp/revoke", payload)

    async def otp_validate(self, otpcode: str) -> dict:
        payload = {"otpcode": otpcode}
        return await self._post("/v1/otp/validate", payload)

    # Policy Management
    async def policy_add_path_exclusions(self, groupid: str, paths: List[str]) -> dict:
        payload = {"groupid": groupid, "path": paths}
        return await self._post("/v1/group/path/add", payload)

    async def policy_add_publishers(self, groupid: str, publishers: List[str]) -> dict:
        payload = {"groupid": groupid, "publisher": publishers}
        return await self._post("/v1/group/publisher/add", payload)

    async def policy_clone(self, source_groupid: str, target_groupid: str) -> dict:
        payload = {"groupid": source_groupid, "targetgroupid": target_groupid}
        return await self._post("/v1/group/assign", payload)

    async def policy_find_all(self) -> pd.DataFrame:
        result = await self._post("/v1/group")
        return pd.DataFrame(result["response"]["groups"])

    async def policy_list_agents(self, groupid: str) -> pd.DataFrame:
        payload = {"groupid": groupid}
        result = await self._post("/v1/group/agents", payload)
        return pd.DataFrame(result["response"]["agents"])

    async def policy_list_allowlists(self, groupid: str) -> pd.DataFrame:
        payload = {"groupid": groupid}
        result = await self._post("/v1/group/policies", payload)
        return pd.DataFrame(result["response"]["applications"])

    async def policy_set_auditmode(self, groupid: str, auditmode: str) -> dict:
        payload = {"groupid": groupid, "auditmode": auditmode}
        return await self._post("/v1/group/settings/auditmode", payload)

    async def policy_set_script_custom(
        self,
        groupid: str,
        script_custom: int,
        scripts_audit: Optional[List[str]] = None,
        scripts_disabled: Optional[List[str]] = None,
        scripts_respect: Optional[List[str]] = None
    ) -> dict:
        payload = {
            "groupid": groupid,
            "script_custom": script_custom,
            "scripts_audit": scripts_audit or [],
            "scripts_disabled": scripts_disabled or [],
            "scripts_respect": scripts_respect or []
        }
        return await self._post("/v1/group/settings/script_custom", payload)

    # Execution History
    async def history_logging(self, type: List[str], checkpoint: str, policy: List[str]) -> List[Dict]:
        payload = {"type": type, "checkpoint": checkpoint, "policy": policy}
        result = await self._post("/v1/logging/exechistories", payload)
        return result["response"]["exechistories"]

    async def history_execution(self, today: str, date_selected: str, agent_name: str) -> List[Dict]:
        payload = {"datefrom": date_selected, "dateto": today, "hostname": agent_name}
        result = await self._post("/v1/getexechistory", payload)
        return result["response"]["exechistory"]