from __future__ import annotations

import abc
import requests
import json
import logging
import jwt
import pkg_resources
from datetime import datetime
from concurrent.futures import Future
from typing import Union, Any, Optional, List, Dict
from pydantic import ValidationError, AnyUrl, PydanticUserError

from airadar.utils import get_git_info
from airadar.utils.config import airadar_configs, ConfigKeys
from airadar.utils.auth import get_ccf_access_token, get_private_key_access_token
from airadar.typing import RadarResponse
from airadar.tracking.model import (
    PipelineRunSchema,
    PipelineRunCreateRequestSchema,
    DatasetVersionCreateRequestSchema,
    DatasetSchema,
    ModelVersionCreateRequestSchema,
    ModelSchema,
    Dependency,
    File,
    PipelineSchema,
    PipelineCreateRequestSchema,
    PipelineRunState,
    PipelineRunUpdateRequestSchema,
    DeploymentCreateRequestSchema,
    DeploymentWithApplicationSchema,
    ExternalRunUrlItem,
    ApplicationUrlItem,
    BuildUrlItem,
    ExternalRunIdItem,
    DescriptionItem,
    RadarUri,
    VersionItem3,
    VersionItem,
    ExternalUrlItem,
)

logger = logging.getLogger(__name__)


RADAR_SDK_VERSION = pkg_resources.get_distribution("airadar").version


class MissingAuthException(Exception):
    """Raised when the environment variables for obtaining access tokens are not set properly"""


class InvalidAuthTokenException(Exception):
    """Raised when we fail to obtain an access token from the auth provider"""


class AccessToken:
    def __init__(self, access_token: str):
        try:
            decoded_token = jwt.decode(
                access_token, options={"verify_signature": False}
            )
            self.expiration_timestamp = decoded_token.get("exp")
        except jwt.DecodeError:
            logger.critical("Failed to decode JWT access token")
            raise InvalidAuthTokenException("Access token is not valid")

        self.access_token = access_token

    def is_expired(self) -> bool:
        if not self.access_token or not self.expiration_timestamp:
            raise InvalidAuthTokenException("Access token is not valid")

        expiration_datetime = datetime.utcfromtimestamp(self.expiration_timestamp)
        return datetime.utcnow() > expiration_datetime


class HTTPClient(metaclass=abc.ABCMeta):
    def __init__(self) -> None:
        """
        Base class for api clients
        """
        self.base_url = str(airadar_configs.get_value(ConfigKeys.AIRADAR_API_SERVER))
        self.api_version_str = (
            airadar_configs.get_value(ConfigKeys.AIRADAR_API_VERSION) or "v1"
        )
        self.base_url = self.base_url.rstrip("/")

        self.headers = {
            "Content-Type": "application/json",
            "User-Agent": f"radar-sdk-python/{RADAR_SDK_VERSION}",
        }

    @abc.abstractmethod
    def post(self, end_point: str, json_data: str) -> Union[Future[Any], RadarResponse]:
        """
        Abstract method for a HTTP POST call. Derived classes should implement this method.

        Arguments:
            end_point: End-point at the base_url
            json_data: JSON payload

        Returns:
            Either an [APIResponse][airadar.api.api_client.APIResponse] or a Future in case of async request.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def patch(
        self, end_point: str, json_data: str
    ) -> Union[Future[Any], RadarResponse]:
        """
        Abstract method for a HTTP PATCH call. Derived classes should implement this method.

        Arguments:
            end_point: End-point at the base_url
            json_data: JSON payload

        Returns:
            Either an [APIResponse][airadar.api.api_client.APIResponse] or a Future in case of async request.
        """
        raise NotImplementedError

    @abc.abstractmethod
    def get(self, end_point: str) -> Union[Future[Any], RadarResponse]:
        """
        Abstract method for a HTTP GET call. Derived classes should implement this method.

        Arguments:
            end_point: End-point at the base_url

        Returns:
            Either an [APIResponse][airadar.api.api_client.APIResponse] or a Future in case of async request.
        """
        raise NotImplementedError


class SyncHTTPClient(HTTPClient):
    def __init__(self) -> None:
        super().__init__()
        self.jwt_access_token: AccessToken = SyncHTTPClient.fetch_access_token()

    @staticmethod
    def fetch_access_token() -> AccessToken:
        auth_client_url = airadar_configs.get_value(ConfigKeys.AIRADAR_AUTH_TENANT_URL)
        auth_client_id = airadar_configs.get_value(ConfigKeys.AIRADAR_AUTH_CLIENT_ID)
        auth_client_secret = airadar_configs.get_value(
            ConfigKeys.AIRADAR_AUTH_CLIENT_SECRET
        )
        auth_private_key = airadar_configs.get_value(
            ConfigKeys.AIRADAR_AUTH_CLIENT_PRIVATE_KEY
        )
        auth_key_id = airadar_configs.get_value(ConfigKeys.AIRADAR_AUTH_CLIENT_KID)
        auth_client_scope = airadar_configs.get_value(
            ConfigKeys.AIRADAR_AUTH_CLIENT_SCOPE
        )

        if (
            auth_client_url is None
            or auth_client_id is None
            or auth_client_scope is None
        ):
            raise MissingAuthException("Auth configurations are missing")

        # Prefer private key auth, otherwise fallback to client credential flow
        if auth_private_key is not None and auth_key_id is not None:
            logger.debug("Using private key auth strategy")
            access_token = get_private_key_access_token(
                str(auth_client_url),
                str(auth_client_id),
                str(auth_private_key),
                str(auth_key_id),
                str(auth_client_scope),
            )
        elif auth_client_secret is not None:
            logger.debug("Using client credential flow auth strategy")
            access_token = get_ccf_access_token(
                str(auth_client_url),
                str(auth_client_id),
                str(auth_client_secret),
                str(auth_client_scope),
            )
        else:
            raise MissingAuthException("Auth configurations are missing")

        if access_token is None:
            logger.critical(
                f"""
                Failed to obtain an access token. Check the logs for more information. 
                Make sure env variables are properly set.
                """
            )
            raise InvalidAuthTokenException("Failed to obtain an access token")

        logger.debug("Obtained a valid access token from auth provider")
        return AccessToken(access_token)

    def _handle_http_method(
        self, method: str, end_point: str, json_data: Optional[str] = None
    ) -> RadarResponse:
        try:
            logger.debug("Checking for a valid access token")
            if not self.jwt_access_token:
                self.jwt_access_token = SyncHTTPClient.fetch_access_token()
            elif self.jwt_access_token.is_expired():
                self.jwt_access_token = SyncHTTPClient.fetch_access_token()
        except (InvalidAuthTokenException, MissingAuthException) as e:
            error_msg = f"Failed to make {method} call due to invalid access token: {e}"
            logger.warning(error_msg)
            return RadarResponse(error_msg=error_msg)

        end_point = end_point.lstrip("/")
        url = f"{self.base_url}/{self.api_version_str}/{end_point}"
        logger.debug("Initiating %s request to endpoint %s", method, url)

        # urllib allows opening both http:// and file:// URLs. To avoid the security risk of
        # accidentally opening the latter, we need to validate the URL is an HTTP endpoint
        if not url.lower().startswith("http"):
            error_msg = f"Invalid URL {url} passed to post_json, skipping request"
            logger.warning(error_msg)
            return RadarResponse(error_msg=error_msg)

        func = None
        try:
            func = getattr(requests, method)
        except AttributeError as ae:
            error_msg = f"HTTP requested method is not part of requests library. {ae}"
            logger.exception(error_msg)
            return RadarResponse(error_msg=error_msg)

        headers = {
            **self.headers,
            "Authorization": f"Bearer {self.jwt_access_token.access_token}",
        }
        try:
            resp = None
            if json_data:
                resp = func(
                    url=url,
                    json=json.loads(json_data),
                    headers=headers,
                    timeout=10,
                )
            else:
                resp = func(url=url, headers=self.headers, timeout=10)

            resp.raise_for_status()
            resp_json = resp.json()
            entity_uuid = resp_json.get("uuid")
            if entity_uuid is None:
                error_msg = f"Could not get entity UUID from response body: {resp_json}"
                logger.error(error_msg)
                return RadarResponse(
                    http_status_code=resp.status_code, error_msg=error_msg
                )

            return RadarResponse(
                http_status_code=resp.status_code,
                entity_uuid=entity_uuid,
                json_response=dict(resp_json),
            )
        except json.decoder.JSONDecodeError as e:
            error_msg = f"Error parsing json payload for request to {url}: {str(e)}"
            logger.error(error_msg)
            return RadarResponse(error_msg=error_msg)
        except requests.exceptions.HTTPError as e:
            if e.response is None:
                return RadarResponse(error_msg="Unknown HTTP Error occurred")
            error_msg = str(
                f"Error with request: {e.response.status_code}\n{e.response.text}"
            )
            logger.error(error_msg)
            return RadarResponse(
                http_status_code=e.response.status_code, error_msg=error_msg
            )
        except requests.exceptions.JSONDecodeError as e:
            error_msg = f"Response was not valid JSON for request to {url}: {str(e)}"
            logger.error(error_msg)
            return RadarResponse(http_status_code=500, error_msg=error_msg)
        except requests.exceptions.Timeout as e:
            error_msg = f"Request timed out for request to {url}"
            logger.error(error_msg)
            return RadarResponse(http_status_code=408, error_msg=error_msg)
        except Exception as e:
            error_msg = f"Unexpected error during request to {url}: {str(e)}"
            logger.error(error_msg)
            status_code = None
            if (
                hasattr(e, "response")
                and hasattr(e.response, "status_code")
                and isinstance(e.response.status_code, int)
            ):
                status_code = e.response.status_code
            return RadarResponse(http_status_code=status_code, error_msg=error_msg)

    def post(self, end_point: str, json_data: str) -> RadarResponse:
        return self._handle_http_method("post", end_point, json_data)

    def patch(self, end_point: str, json_data: str) -> RadarResponse:
        return self._handle_http_method("patch", end_point, json_data)

    def get(self, end_point: str) -> RadarResponse:
        return self._handle_http_method("get", end_point)


class RadarAPIClient:
    def __init__(self) -> None:
        self.http_client = self._setup_http_client()

    def _setup_http_client(self) -> Union[SyncHTTPClient, None]:
        try:
            return SyncHTTPClient()
        except (InvalidAuthTokenException, MissingAuthException):
            pass
        except Exception as e:
            logger.critical(
                "Unexpected error when fetching Radar access token: %s", str(e)
            )

        return None

    def is_authorized(self) -> bool:
        return self.http_client is not None

    def _validate_api_response(
        self,
        response: RadarResponse,
        model_class: Any,
    ) -> RadarResponse:
        try:
            model_class.model_validate(response.json_response)
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=response.http_status_code,
                error_msg=error_msg,
                json_response=response.json_response,
            )

        return response

    def post_pipeline(
        self, pipeline_name: str, external_url: Optional[ExternalUrlItem] = None
    ) -> RadarResponse:
        if self.http_client is None:
            return RadarResponse(error_msg="Client data collection disabled")

        try:
            pipeline_create_obj = PipelineCreateRequestSchema(
                name=pipeline_name, external_url=external_url
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Failed to pipeline create request object."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        pipeline_create_res = self.http_client.post(
            "/pipelines",
            pipeline_create_obj.model_dump_json(),
        )

        if not pipeline_create_res.is_success():
            logger.critical(
                "AI Radar pipeline could not be created, no data will be collected by AI Radar for this run"
            )
            return pipeline_create_res

        try:
            pipeline = PipelineSchema.model_validate(pipeline_create_res.json_response)
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for pipeline create request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            self.disable_data_collection = True
            return RadarResponse(
                http_status_code=pipeline_create_res.http_status_code,
                error_msg=error_msg,
                json_response=pipeline_create_res.json_response,
            )

        # Set the pipeline run id here to avoid duplicate calls when airadar.log_artifact() is called
        # following a call to airadar.active_run()
        self.pipeline_id = pipeline.uuid
        self.pipeline = pipeline
        return pipeline_create_res

    def post_pipeline_run(
        self,
        pipeline_id: str,
        external_run_id: Optional[str] = None,
        run_at: Optional[datetime] = None,
        run_uri: Optional[str] = None,
        dependency_report: List[Dependency] = [],
    ) -> RadarResponse:
        if self.http_client is None:
            return RadarResponse(error_msg="Client data collection disabled")

        git_info = get_git_info()

        try:
            # TODO: Populate platform with a meaningful name
            # See https://github.com/protectai/radar-sdk/issues/72
            create_run_obj = PipelineRunCreateRequestSchema(
                external_run_id=(
                    ExternalRunIdItem(root=external_run_id) if external_run_id else None
                ),
                external_run_url=(
                    ExternalRunUrlItem(root=AnyUrl(run_uri)) if run_uri else None
                ),
                dependencies=dependency_report,
                run_at=run_at or datetime.now(),
                git_info=git_info,
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Error when trying to create request based on given inputs."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        run_create_res = self.http_client.post(
            f"/pipelines/{pipeline_id}/runs",
            create_run_obj.model_dump_json(),
        )

        if not run_create_res.is_success():
            logger.critical(
                "AI Radar pipeline run ID could not be retrieved, no data will be collected by AI Radar for this run"
            )
            return run_create_res

        try:
            pipeline_run = PipelineRunSchema.model_validate(
                run_create_res.json_response
            )
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for pipeline run create request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            self.disable_data_collection = True
            return RadarResponse(
                http_status_code=run_create_res.http_status_code,
                error_msg=error_msg,
                json_response=run_create_res.json_response,
            )

        # Set the pipeline run id here to avoid duplicate calls when airadar.log_artifact() is called
        # following a call to airadar.active_run()
        self.pipeline_run_id = pipeline_run.uuid
        self.pipeline_run = pipeline_run
        return run_create_res

    def patch_pipeline_run(
        self,
        status: PipelineRunState,
        pipeline_run_id: str,
        completed_at: Optional[datetime] = None,
    ) -> RadarResponse:
        if self.http_client is None:
            return RadarResponse(error_msg="Client data collection disabled")

        try:
            patch_run_obj = PipelineRunUpdateRequestSchema(
                status=status, completed_at=completed_at
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Error when trying to create request based on given inputs."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        run_patch_res = self.http_client.patch(
            f"/pipelines/run/{pipeline_run_id}",
            patch_run_obj.model_dump_json(),
        )

        if not run_patch_res.is_success():
            logger.critical(
                f"Failed to mark pipeline run {pipeline_run_id} complete in AI Radar"
            )
            return run_patch_res

        try:
            patch_response = PipelineRunSchema.model_validate(
                run_patch_res.json_response
            )
            logger.info(
                f"Pipeline run with uuid={patch_response.uuid} updated successfully"
            )
            return run_patch_res
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for pipeline run patch request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=run_patch_res.http_status_code,
                error_msg=error_msg,
                json_response=run_patch_res.json_response,
            )

    def post_model(
        self,
        name: str,
        uri: str,
        files: List[File],
        version: Optional[str] = None,
        pipeline_run_id: Optional[str] = None,
        description: Optional[str] = None,
        properties: Optional[Dict[str, str]] = None,
    ) -> RadarResponse:
        if not self.http_client:
            return RadarResponse(error_msg="Client data collection disabled")

        if pipeline_run_id:
            model_version_endpoint = f"/pipelines/run/{pipeline_run_id}/model-versions"
        else:
            model_version_endpoint = "/models"

        try:
            request = ModelVersionCreateRequestSchema(
                name=name,
                external_uri=uri,
                version=VersionItem3(root=version) if version else None,
                description=DescriptionItem(root=description) if description else None,
                files=files,
                properties=properties,
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Error when trying to create request based on given inputs."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        response = self.http_client.post(
            model_version_endpoint, request.model_dump_json(by_alias=True)
        )

        if not response.is_success():
            logger.critical(f"Failed to create model version for {name}")
            return response

        try:
            model_response = ModelSchema.model_validate(response.json_response)
            logger.info(f"Model version created with uuid={model_response.uuid}")
            return response
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for model version create request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=response.http_status_code,
                error_msg=error_msg,
                json_response=response.json_response,
            )

    def post_input_model(self, radar_uri: str, pipeline_run_id: str) -> RadarResponse:
        if not self.http_client:
            return RadarResponse(error_msg="Client data collection disabled")

        associate_model_endpoint = f"/pipelines/run/{pipeline_run_id}/associate-model"

        response = self.http_client.post(
            associate_model_endpoint, json.dumps({"radar_uri": radar_uri})
        )

        if not response.is_success():
            logger.critical(
                f"Failed to associate input model version at {radar_uri} with pipeline run {pipeline_run_id}"
            )
            return response

        try:
            model_response = ModelSchema.model_validate(response.json_response)
            if model_response.latest_version:
                logger.info(
                    f"Model version uuid={model_response.latest_version.uuid} was successfully associated with pipeline run {pipeline_run_id}"
                )
            return response
        except ValidationError as e:
            error_msg = f"Got an invalid response from the API server when trying to associate input model {radar_uri}"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=response.http_status_code,
                error_msg=error_msg,
                json_response=response.json_response,
            )

    def post_dataset(
        self,
        name: str,
        uri: str,
        files: List[File],
        pipeline_run_id: str,
        version: Optional[str] = None,
        description: Optional[str] = None,
    ) -> RadarResponse:
        if not self.http_client:
            return RadarResponse(error_msg="Client data collection disabled")

        dataset_version_endpoint = f"/pipelines/run/{pipeline_run_id}/dataset-versions"

        try:
            request = DatasetVersionCreateRequestSchema(
                name=name,
                external_uri=uri,
                version=VersionItem(root=version) if version else None,
                description=DescriptionItem(root=description) if description else None,
                files=files,
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Error when trying to create request based on given inputs."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        response = self.http_client.post(
            dataset_version_endpoint, request.model_dump_json()
        )

        if not response:
            logger.critical(f"Failed to create dataset version for {name}")
            return response

        try:
            dataset_response = DatasetSchema.model_validate(response.json_response)
            logger.info(f"Dataset version created with uuid={dataset_response.uuid}")
            return response
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for dataset version create request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=response.http_status_code,
                error_msg=error_msg,
                json_response=response.json_response,
            )

    def post_deployment(
        self,
        application_id: str,
        application_version: str,
        environment: str,
        application_url: Optional[str] = None,
        build_url: Optional[str] = None,
        radar_uris: Optional[List[str]] = None,
        build_time: datetime = datetime.now(),
    ) -> RadarResponse:
        if not self.http_client:
            return RadarResponse(error_msg="Client data collection disabled")

        try:
            deployment = DeploymentCreateRequestSchema(
                application_version=application_version,
                environment=environment,
                application_url=(
                    ApplicationUrlItem(root=AnyUrl(application_url))
                    if application_url
                    else None
                ),
                build_url=BuildUrlItem(root=AnyUrl(build_url)) if build_url else None,
                build_time=build_time,
                radar_uris=(
                    [RadarUri(root=radar_uri) for radar_uri in radar_uris]
                    if radar_uris
                    else []
                ),
            )
        except (ValidationError, PydanticUserError) as e:
            error_msg = "Error when trying to create request based on given inputs."
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(error_msg=error_msg)

        response = self.http_client.post(
            f"/applications/{application_id}/deployments",
            deployment.model_dump_json(),
        )
        if not response.is_success():
            logger.critical(
                f"Failed to create a deployment for application id {application_id}"
            )
            return response

        try:
            deployment_response = DeploymentWithApplicationSchema.model_validate(
                response.json_response
            )
            logger.info(f"Deployment created with uuid={deployment_response.uuid}")
            return response
        except ValidationError as e:
            error_msg = "Got an invalid response from the API server for deployment create request"
            logger.critical(
                error_msg,
                e,
                exc_info=True,
            )
            return RadarResponse(
                http_status_code=response.http_status_code,
                error_msg=error_msg,
                json_response=response.json_response,
            )
