import io
import json
import logging
import os
import uuid
from urllib.parse import urlencode, urlunparse

import numpy as np
import requests
from PIL import Image
from relai_visualization.models.application_result import ApplicationResult

RELAI_API_KEY = "RELAI_API_KEY"
RELAI_PERMISSION_GROUP = "RELAI_PERMISSION_GROUP"
RELAI_API_MODE = os.environ.get("RELAI_API_MODE", "production")
if RELAI_API_MODE == "production":
    RELAI_API_URL = "api.relai.ai"
    RELAI_API_PROTOCOL = "https"
else:
    RELAI_API_URL = "localhost:8889"
    RELAI_API_PROTOCOL = "http"
RELAI_API_VERSION = "v1"


class RELAIClient:
    """Client for interacting with the RELAI API."""

    schema_endpoint: str = urlunparse((RELAI_API_PROTOCOL, RELAI_API_URL, "schema/", "", "", ""))
    run_id_endpoint: str = urlunparse(
        (RELAI_API_PROTOCOL, RELAI_API_URL, f"api/{RELAI_API_VERSION}/relai-run/", "", "", "")
    )
    llm_serve_endpoint: str = urlunparse(
        (RELAI_API_PROTOCOL, RELAI_API_URL, f"api/{RELAI_API_VERSION}/relai-llm-serve/", "", "", "")
    )
    web_search_endpoint: str = urlunparse(
        (RELAI_API_PROTOCOL, RELAI_API_URL, f"api/{RELAI_API_VERSION}/relai-web-search/", "", "", "")
    )

    @classmethod
    def credit_balance_endpoint(cls, app_name: str, num_samples: int) -> str:
        """Generates the URL for the credit balance endpoint.

        Args:
            app_name (str): The name of the application.
            num_samples (int): The number of sample units for which the potential credits usage is to be estimated.

        Returns:
            str: The URL for the credit balance endpoint.
        """
        query_params = {
            "app_name": app_name,
            "number_units": num_samples,
        }
        return urlunparse(
            (
                RELAI_API_PROTOCOL,
                RELAI_API_URL,
                f"api/{RELAI_API_VERSION}/relai-license-details/remaining-credits/",
                "",
                urlencode(query_params),
                "",
            )
        )

    @classmethod
    def analysis_result_endpoint(cls, run_id: str) -> str:
        """Generates the URL for uploading/downloading application results.

        Args:
            run_id (str): The run id to which the result should be logged.

        Returns:
            str: The URL for uploading application results.
        """
        return urlunparse(
            (
                RELAI_API_PROTOCOL,
                RELAI_API_URL,
                f"api/{RELAI_API_VERSION}/relai-run/{run_id}/analysis_result/",
                "",
                "",
                "",
            )
        )

    @classmethod
    def transaction_endpoint(cls, app_name: str, num_samples_processed: int, run_id: str) -> str:
        """Generates the URL for logging credit transactions.

        Args:
            app_name (str): The name of the application.
            num_samples_processed (int): The number of sample units processed in this analysis.
            run_id (str): The run id to which the result should be logged.

        Returns:
            str: The URL for logging credit transactions.
        """
        query_params = {
            "app_name": app_name,
            "number_units": num_samples_processed,
            "run_id": run_id,
        }
        return urlunparse(
            (
                RELAI_API_PROTOCOL,
                RELAI_API_URL,
                f"api/{RELAI_API_VERSION}/relai-credit-transactions/",
                "",
                urlencode(query_params),
                "",
            )
        )

    def __init__(self) -> None:
        # Retrieve RELAI API keys
        self._relai_api_key = os.getenv(RELAI_API_KEY)
        self._relai_permission_group = os.getenv(RELAI_PERMISSION_GROUP)
        if self._relai_api_key is None:
            raise ValueError(f"{RELAI_API_KEY} environment variable not set. Please set it to with your API key.")
        if self._relai_permission_group is None:
            raise ValueError(
                f"{RELAI_PERMISSION_GROUP} environment variable not set. Please set it to with your RELAI permission group."
            )

    def get_run_id(self, project_name: str, run_name: str) -> str:
        """Retrieves an id for an application run.

        Args:
            project_name (str): Name of the project.
            run_name (str): Name of the run.

        Returns:
            str: The run id.
        """
        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Relai-Permission-Group": f"{self._relai_permission_group}",
            "Content-Type": "application/json",
        }
        data = {
            "project_id": project_name,
            "name": run_name,
        }
        data = json.dumps(data)
        response = requests.post(self.run_id_endpoint, headers=headers, data=data)
        if response.status_code != 201:
            error_message = response.json().get("detail", "")
            raise RuntimeError(f"Run Id could not be generated. {error_message}")
        run_id = response.json()["run_id"]
        return run_id

    def has_sufficient_credits(self, app_name: str, num_samples: int) -> tuple[bool, float, float]:
        """Retrieves the current balance of RELAI credits available under the license and estimates the potential
        credits usage for a given application and given number of sample units. Designed to help manage and
        forecast credit consumption for processing tasks.

        Args:
            app_name (str): The name of the application.
            num_samples (int): The number of sample units for which the potential credits usage is to be estimated.

        Returns:
            tuple[bool, float, float]: A tuple containing three elements:
                - A boolean indicating if there are enough RELAI credits or not.
                - The remaining credits available under the license.
                - The estimated potential credits usage for the specified number of sample units.
        """
        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Relai-Permission-Group": f"{self._relai_permission_group}",
            "Content-Type": "application/json",
        }
        try:
            response = requests.get(
                RELAIClient.credit_balance_endpoint(app_name, num_samples),
                headers=headers,
            )
            if response.status_code == 200:
                data = response.json()
                remaining_credits = data.get("remaining_credits", np.nan)
                potential_credits_usage = data.get("potential_credits_usage", np.nan)
                return True, remaining_credits, potential_credits_usage
            elif response.status_code == 402:
                data = response.json()
                remaining_credits = data.get("remaining_credits", np.nan)
                potential_credits_usage = data.get("potential_credits_usage", np.nan)
                return False, remaining_credits, potential_credits_usage
            else:
                error_message = response.json().get("detail", "")
                logging.exception(f"Error while fetching RELAI credits. {error_message}")
                # return False if there is an error
                return False, np.nan, np.nan
        except Exception as e:
            logging.exception(f"Error while fetching RELAI credits. {e}.")
            # return False if there is an error
            return False, np.nan, np.nan

    def log_metadata(
        self,
        metadata: dict,
        run_id: str,
    ):
        """Logs `metadata` for this RELAI run. To be displayed on RELAI dashboard.

        NOTE: Metadata should be logged before results. Calling this after log_results over-writes the results.

        Args:
            metadata (dict): The metadata for this RELAI run.
            run_id (str): The run id to which the result should be logged.
        """

        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Content-Type": "application/json",
        }
        try:
            response = requests.request(
                "POST",
                RELAIClient.analysis_result_endpoint(run_id),
                headers=headers,
                data=json.dumps(metadata),
            )
            if response.status_code != 200:
                error_message = response.json().get("detail", "")
                raise RuntimeError(f"Result data could not be logged. {error_message}")
        except Exception as e:
            logging.exception(f"Error while logging the result. {e}.")
            raise e

    def log_result(
        self,
        application_result: ApplicationResult,
        run_id: str,
        num_samples_processed: int,
    ):
        """Logs `ApplicationResult` for visualization on the RELAI dashboard.

        Args:
            application_result (ApplicationResult): The application result to be logged.
            run_id (str): The run id to which the result should be logged.
            num_samples_processed (int): The number of sample units processed in this analysis.
        """
        # Log the credit transaction
        app_name = application_result.data["application"]
        try:
            credit_headers = {
                "Authorization": f"Token {self._relai_api_key}",
                "Relai-Permission-Group": f"{self._relai_permission_group}",
                "Content-Type": "application/json",
            }
            response = requests.request(
                "POST",
                RELAIClient.transaction_endpoint(app_name, num_samples_processed, run_id),
                headers=credit_headers,
            )
            if response.status_code != 200:
                error_message = response.json().get("detail", "")
                raise RuntimeError(f"Error: Could not complete RELAI credits transaction. {error_message}")
        except Exception as e:
            error_message = f"Error: Could not complete RELAI credits transaction. {e}."
            logging.exception(error_message)
            raise RuntimeError(error_message)

        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Content-Type": "application/json",
        }
        data = application_result.json()
        try:
            response = requests.request(
                "POST",
                RELAIClient.analysis_result_endpoint(run_id),
                headers=headers,
                data=data,
            )
            if response.status_code != 200:
                error_message = response.json().get("detail", "")
                raise RuntimeError(f"Result data could not be logged. {error_message}")
        except Exception as e:
            logging.exception(f"Error while logging the result. {e}.")
            raise e

    def upload_image(self, image: Image.Image, run_id: str) -> str:
        """Logs image for visualization on the RELAI dashboard and returns the image URL.

        Args:
            image (Image.Image): A PIL image to be uploaded.
            run_id (str): The run id to which the image should be logged.

        Returns:
            str: URL of the uploaded image.
        """
        headers = {
            "Authorization": f"Token {self._relai_api_key}",
        }
        if not isinstance(image, Image.Image):
            raise ValueError("Image must be a PIL image object.")

        try:
            buffer = io.BytesIO()
            image.save(buffer, format="PNG")
            image_type = "image/png"
            image_identifier = f"{uuid.uuid4().hex}.png"

            files = [
                (image_identifier, (image_identifier, buffer.getvalue(), image_type)),
            ]
            response = requests.request(
                "POST",
                RELAIClient.analysis_result_endpoint(run_id),
                headers=headers,
                files=files,
            )
        except Exception as e:
            logging.exception(f"Error while logging the image. {e}.")
            raise e

        if response.status_code != 200:
            error_message = response.json().get("detail", "")
            raise RuntimeError(f"Image could not be logged. {error_message}")
        else:
            image_url = response.json().get("image_url", "")
            return image_url

    def get_result(self, run_id: str):
        """Get `ApplicationResult` for a particular RELAI run.

        Args:
            run_id (str): The run id whose `ApplicationResult` we want to get.
        """
        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Content-Type": "application/json",
        }
        params = {"data": "visualization_data"}
        try:
            response = requests.request(
                "GET",
                RELAIClient.analysis_result_endpoint(run_id),
                headers=headers,
                params=params,
            )
            if response.status_code != 200:
                error_message = response.json().get("detail", "")
                raise RuntimeError(f"Result data could not be retrieved. {error_message}")

            response_json = response.json()
            # Check if the run is not complete
            if "detail" in response_json.keys() and "not completed" in response_json["detail"]:
                raise RuntimeError("Run is not complete")
            return response.json()["visualization_data"]
        except Exception as e:
            logging.exception(f"Error while retrieving the result. {e}.")
            raise e

    def get_search_results(self, query: str, num_results: int = 10) -> list[dict]:
        """Get web search results for a given query.

        Args:
            query (str): The search query.
            num_results (int): The number of search results to retrieve.

        Returns:
            list[dict]: The search results.
        """
        headers = {
            "Authorization": f"Token {self._relai_api_key}",
            "Content-Type": "application/json",
        }
        params = {"query": query, "num_results": num_results}
        try:
            response = requests.request(
                "GET",
                RELAIClient.web_search_endpoint,
                headers=headers,
                params=params,
            )
            if response.status_code != 200:
                error_message = response.json().get("detail", "")
                raise RuntimeError(f"Search results could not be retrieved. {error_message}")

            return response.json()
        except Exception as e:
            logging.exception(f"Error while retrieving search results. {e}.")
            raise e
