import gzip
import json
import logging
from typing import Optional
import httpx
import msgpack
from acelerai_pipelines import APIGW, VERIFY_HTTPS
from acelerai_pipelines.models.deploy import Deploy
from acelerai_pipelines.models.run import Run
from acelerai_pipelines.models.pagination_data import PaginationData
from acelerai_inputstream.utils import CustomJSONEncoder, CustomJsonDecoder, custom_encoder

logger = logging.getLogger("HttpAcelerAI")
packer = msgpack.Packer(default=custom_encoder)


class BaseAcelerAIHttpPipelinesClient:
    """
    Base HTTP client with common headers and helpers for Aceler Pipelines.
    """

    def __init__(self, token: str):
        self.token = token

    def get_headers(self) -> dict:
        """
        Get the standard headers for requests.
        """
        return {
            "Authorization": f"A2G {self.token}",
            "Content-Type": "application/json"
        }

    def _prepare_process_status_payload(
        self,
        run_id: str,
        status: str,
        process_name: str,
        title: str,
        content: Optional[str] = None
    ) -> tuple[dict, bytes]:
        """
        Builds headers with gzip encoding and compresses the JSON body.
        """
        body = {
            "runId": run_id,
            "processName": process_name,
            "status": status,
            "title": title,
            "content": content
        }
        headers = self.get_headers()
        headers["Content-Encoding"] = "gzip"

        serialized = json.dumps(body, cls=CustomJSONEncoder).encode("utf-8")
        compressed = gzip.compress(serialized)
        return headers, compressed

    def _prepare_execute_payload(self, deploy_key: str, payload: Optional[dict] = None) -> dict:
        return {
            "fKey": deploy_key,
            "payload": json.dumps(payload, cls=CustomJSONEncoder) if payload is not None else None
        }


class AcelerAIHttpPipelinesClient(BaseAcelerAIHttpPipelinesClient):
    """
    Synchronous HTTP client for Aceler Pipelines API.
    """

    def __init__(self, token: str):
        super().__init__(token)
        self._client: Optional[httpx.Client] = None

    def _get_client(self) -> httpx.Client:
        if self._client is None or self._client.is_closed:
            self._client = httpx.Client(
                verify=VERIFY_HTTPS,
                timeout=httpx.Timeout(10800.0, connect=10.0, read=10800.0)
            )
        return self._client

    def close(self):
        """Close the underlying HTTP client session."""
        if self._client is not None and not self._client.is_closed:
            self._client.close()

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()

    def get_deploy_by_fkey(self, fkey: str) -> Optional[Deploy]:
        """
        Get deploy by fkey.
        """
        try:
            client = self._get_client()
            response = client.get(f"{APIGW}/Deploy/Fkey/{fkey}", headers=self.get_headers())
            logger.debug(f"response {response.status_code} - {response.text}")
            if response.status_code == 200:
                resp_json = response.json(cls=CustomJsonDecoder)
                data = resp_json["data"]
                return Deploy(**data)
            else:
                logger.error(f"Error getting deploy by fkey: {fkey}")
                return None
        except Exception as e:
            logger.error(f"Error getting deploy by fkey: {fkey}", exc_info=True)
            return None

    def execute_run_ondemand(self, deploy_key: str, payload: Optional[dict] = None) -> Optional[Run]:
        """
        Execute a run on demand.
        """
        try:
            body = self._prepare_execute_payload(deploy_key, payload)
            client = self._get_client()
            response = client.post(f"{APIGW}/Run/ExecuteOnDemand", headers=self.get_headers(), json=body)
            if response.status_code == 201:
                json_res = response.json(cls=CustomJsonDecoder)
                data = json_res["data"]
                return Run(**data)
            else:
                logger.error(f"Error executing run on demand: {response.status_code} - {response.text}")
                return None
        except Exception as e:
            logger.error(f"Error executing run on demand - Deploy ID: {deploy_key}", exc_info=True)
            return None

    def get_run(self, run_id: str) -> Optional[Run]:
        """
        Get run by id.
        """
        try:
            client = self._get_client()
            response = client.get(f"{APIGW}/Run/{run_id}", headers=self.get_headers())
            if response.status_code == 200:
                data = response.json()
                return Run(**data["data"])
            else:
                logger.error(f"Error getting run - Run ID: {run_id} - {response.text}")
                return None
        except Exception as e:
            logger.error(f"Error getting run - Run ID: {run_id}", exc_info=True)
            return None

    def get_deploy_runs(self, deploy_key: str, page: int = 1, page_size: int = 10) -> Optional[PaginationData[Run]]:
        """
        Get paginated runs for a deploy.
        """
        try:
            client = self._get_client()
            response = client.get(
                f"{APIGW}/Deploy/{deploy_key}/Run/History?Page={page}&PageSize={page_size}",
                headers=self.get_headers()
            )
            if response.status_code == 200:
                data = response.json()
                runs = [Run(**run) for run in data["data"]["data"]]
                total = data["data"]["total"]
                return PaginationData(runs, page, page_size, total)
            elif response.status_code == 400:
                logger.error(f"Error getting runs - Deploy Key: {deploy_key} - {response.text}")
            elif response.status_code == 404:
                logger.error(f"Deploy {deploy_key} not found in your account, please check the deploy key or your credentials")
            return None
        except Exception as e:
            logger.error(f"Error getting runs - Deploy Key: {deploy_key}", exc_info=True)
            return None

    def send_process_status(
        self,
        run_id: str,
        status: str,
        process_name: str,
        title: str,
        content: Optional[str] = None
    ):
        """
        Send process status for a run synchronously.
        """
        try:
            headers, compressed = self._prepare_process_status_payload(run_id, status, process_name, title, content)
            client = self._get_client()
            response = client.post(
                f"{APIGW}/StatusControl/ProcessStatus",
                headers=headers,
                content=compressed,
                timeout=httpx.Timeout(30.0, connect=10.0)
            )
            if response.status_code == 200:
                logger.debug(f"ProcessStatus sent successfully for run {run_id}")
            else:
                err_msg = f"Error sending ProcessStatus: {response.status_code} - {response.text}"
                logger.error(err_msg)
                raise Exception(err_msg)
        except Exception as e:
            logger.error(f"Error sending ProcessStatus for run {run_id}", exc_info=True)
            raise


class AcelerAIAsyncHttpPipelinesClient(BaseAcelerAIHttpPipelinesClient):
    """
    Asynchronous HTTP client for Aceler Pipelines API.
    """

    def __init__(self, token: str):
        super().__init__(token)
        self._client: Optional[httpx.AsyncClient] = None

    def _get_client(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(
                verify=VERIFY_HTTPS,
                timeout=httpx.Timeout(10800.0, connect=10.0, read=10800.0)
            )
        return self._client

    async def close(self):
        """Close the underlying async HTTP client session."""
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close()

    async def get_deploy_by_fkey(self, fkey: str) -> Optional[Deploy]:
        """
        Get deploy by fkey asynchronously.
        """
        try:
            client = self._get_client()
            response = await client.get(f"{APIGW}/Deploy/Fkey/{fkey}", headers=self.get_headers())
            logger.debug(f"response {response.status_code} - {response.text}")
            if response.status_code == 200:
                resp_json = response.json(cls=CustomJsonDecoder)
                data = resp_json["data"]
                return Deploy(**data)
            else:
                logger.error(f"Error getting deploy by fkey: {fkey}")
                return None
        except Exception as e:
            logger.error(f"Error getting deploy by fkey: {fkey}", exc_info=True)
            return None

    async def execute_run_ondemand(self, deploy_key: str, payload: Optional[dict] = None) -> Optional[Run]:
        """
        Execute a run on demand asynchronously.
        """
        try:
            body = self._prepare_execute_payload(deploy_key, payload)
            client = self._get_client()
            response = await client.post(f"{APIGW}/Run/ExecuteOnDemand", headers=self.get_headers(), json=body)
            if response.status_code == 201:
                json_res = response.json(cls=CustomJsonDecoder)
                data = json_res["data"]
                return Run(**data)
            else:
                logger.error(f"Error executing run on demand: {response.status_code} - {response.text}")
                return None
        except Exception as e:
            logger.error(f"Error executing run on demand - Deploy ID: {deploy_key}", exc_info=True)
            return None

    async def get_run(self, run_id: str) -> Optional[Run]:
        """
        Get run by id asynchronously.
        """
        try:
            client = self._get_client()
            response = await client.get(f"{APIGW}/Run/{run_id}", headers=self.get_headers())
            if response.status_code == 200:
                data = response.json()
                return Run(**data["data"])
            else:
                logger.error(f"Error getting run - Run ID: {run_id} - {response.text}")
                return None
        except Exception as e:
            logger.error(f"Error getting run - Run ID: {run_id}", exc_info=True)
            return None

    async def get_deploy_runs(self, deploy_key: str, page: int = 1, page_size: int = 10) -> Optional[PaginationData[Run]]:
        """
        Get paginated runs for a deploy asynchronously.
        """
        try:
            client = self._get_client()
            response = await client.get(
                f"{APIGW}/Deploy/{deploy_key}/Run/History?Page={page}&PageSize={page_size}",
                headers=self.get_headers()
            )
            if response.status_code == 200:
                data = response.json()
                runs = [Run(**run) for run in data["data"]["data"]]
                total = data["data"]["total"]
                return PaginationData(runs, page, page_size, total)
            elif response.status_code == 400:
                logger.error(f"Error getting runs - Deploy Key: {deploy_key} - {response.text}")
            elif response.status_code == 404:
                logger.error(f"Deploy {deploy_key} not found in your account, please check the deploy key or your credentials")
            return None
        except Exception as e:
            logger.error(f"Error getting runs - Deploy Key: {deploy_key}", exc_info=True)
            return None

    async def send_process_status(
        self,
        run_id: str,
        status: str,
        process_name: str,
        title: str,
        content: Optional[str] = None
    ):
        """
        Send process status for a run asynchronously.
        """
        try:
            headers, compressed = self._prepare_process_status_payload(run_id, status, process_name, title, content)
            client = self._get_client()
            response = await client.post(
                f"{APIGW}/StatusControl/ProcessStatus",
                headers=headers,
                content=compressed,
                timeout=httpx.Timeout(30.0, connect=10.0)
            )
            if response.status_code == 200:
                logger.debug(f"ProcessStatus sent successfully for run {run_id}")
            else:
                err_msg = f"Error sending ProcessStatus: {response.status_code} - {response.text}"
                logger.error(err_msg)
                raise Exception(err_msg)
        except Exception as e:
            logger.error(f"Error sending ProcessStatus for run {run_id}", exc_info=True)
            raise
