import logging
import math
from typing import Optional, Any, NamedTuple, Union, Tuple
import aiohttp
import requests
from enum import Enum
import os
from tenacity import (
    retry,
    wait_exponential_jitter,
    retry_if_exception_type,
    RetryCallState,
)
from pydantic import SecretStr

import deepeval
from deepeval.key_handler import KEY_FILE_HANDLER, KeyValues
from deepeval.confident.types import ApiResponse, ConfidentApiError
from deepeval.config.settings import get_settings

CONFIDENT_API_KEY_ENV_VAR = "CONFIDENT_API_KEY"
DEEPEVAL_BASE_URL = "https://deepeval.confident-ai.com"
DEEPEVAL_BASE_URL_EU = "https://eu.deepeval.confident-ai.com"
DEEPEVAL_BASE_URL_AU = "https://au.deepeval.confident-ai.com"
API_BASE_URL = "https://api.confident-ai.com"
API_BASE_URL_EU = "https://eu.api.confident-ai.com"
API_BASE_URL_AU = "https://au.api.confident-ai.com"
retryable_exceptions = requests.exceptions.SSLError


def _infer_region_from_api_key(api_key: Optional[str]) -> Optional[str]:
    """
    Infer region from Confident API key prefix.

    Supported:
      - confident_eu_... => "EU"
      - confident_us_... => "US"
      - confident_au_... => "AU"

    Returns None if prefix is not recognized or api_key is falsy.
    """
    if not api_key:
        return None
    key = api_key.strip().lower()
    if key.startswith("confident_eu_"):
        return "EU"
    if key.startswith("confident_us_"):
        return "US"
    if key.startswith("confident_au_"):
        return "AU"
    return None


class BackendResolution(NamedTuple):
    base_url: str
    region: Optional[str]
    # "custom_base_url" | "explicit_region" | "api_key_prefix" | "default"
    source: str


def _api_url_for_region(region: Optional[str]) -> str:
    if region == "EU":
        return API_BASE_URL_EU
    if region == "AU":
        return API_BASE_URL_AU
    return API_BASE_URL


def resolve_backend() -> BackendResolution:
    """Resolve the Confident AI API base URL, along with the region and how
    it was decided.

    Single source of truth for backend routing: `get_base_api_url()` (used
    by all API calls) and `deepeval diagnose` (which displays the region and
    its provenance) both go through here, so the two cannot drift.
    """
    s = get_settings()
    if s.CONFIDENT_BASE_URL:
        return BackendResolution(
            s.CONFIDENT_BASE_URL.rstrip("/"), None, "custom_base_url"
        )

    # If the user has explicitly set a region, respect it.
    region = KEY_FILE_HANDLER.fetch_data(KeyValues.CONFIDENT_REGION)
    if region:
        return BackendResolution(
            _api_url_for_region(region), region, "explicit_region"
        )

    # Otherwise, infer region from the API key prefix.
    inferred = _infer_region_from_api_key(get_confident_api_key())
    if inferred:
        return BackendResolution(
            _api_url_for_region(inferred), inferred, "api_key_prefix"
        )

    # Default to US (backwards compatible)
    return BackendResolution(API_BASE_URL, "US", "default")


def get_base_api_url():
    return resolve_backend().base_url


def get_confident_api_key() -> Optional[str]:
    s = get_settings()
    key: Optional[SecretStr] = s.CONFIDENT_API_KEY
    return key.get_secret_value() if key else None


def set_confident_api_key(api_key: Optional[str]) -> None:
    """
    - Always updates runtime (os.environ) via settings.edit()
    - If DEEPEVAL_DEFAULT_SAVE is set, also persists to dotenv
    - Never writes secrets to the legacy JSON keystore (your Settings logic already skips secrets)
    """
    s = get_settings()
    save = (
        s.DEEPEVAL_DEFAULT_SAVE or None
    )  # e.g. "dotenv" or "dotenv:/path/.env"

    # If you *only* want runtime changes unless a default save is present:
    if save is None:
        with s.edit(persist=False):
            s.CONFIDENT_API_KEY = SecretStr(api_key) if api_key else None
    else:
        # Respect default save: update runtime + write to dotenv, but not JSON
        with s.edit(save=save, persist=None):
            s.CONFIDENT_API_KEY = SecretStr(api_key) if api_key else None


def is_confident():
    return get_confident_api_key() is not None


def log_retry_error(retry_state: RetryCallState):
    exception = retry_state.outcome.exception()
    logging.error(
        f"Confident AI Error: {exception}. Retrying: {retry_state.attempt_number} time(s)..."
    )


class HttpMethods(Enum):
    GET = "GET"
    POST = "POST"
    DELETE = "DELETE"
    PUT = "PUT"


class Endpoints(Enum):
    DATASET_ALIAS_ENDPOINT = "/v1/datasets/:alias"
    DATASET_ALIAS_QUEUE_ENDPOINT = "/v1/datasets/:alias/queue"
    DATASET_ALIAS_VERSIONS_ENDPOINT = "/v1/datasets/:alias/versions"

    TEST_RUN_ENDPOINT = "/v1/test-run"
    EXPERIMENT_ENDPOINT = "/v1/experiment"
    TRACES_ENDPOINT = "/v1/traces"
    ANNOTATIONS_ENDPOINT = "/v1/annotations"
    PROMPTS_VERSION_ID_ENDPOINT = "/v1/prompts/:alias/versions/:version"
    PROMPTS_LABEL_ENDPOINT = "/v1/prompts/:alias/labels/:label"
    PROMPTS_ENDPOINT = "/v1/prompts"
    PROMPTS_VERSIONS_ENDPOINT = "/v1/prompts/:alias/versions"
    PROMPTS_COMMITS_ENDPOINT = "/v1/prompts/:alias/commits"
    PROMPTS_COMMIT_HASH_ENDPOINT = "/v1/prompts/:alias/commits/:hash"
    PROMPTS_BRANCHES_ENDPOINT = "/v1/prompts/:alias/branches"
    PROMPTS_BRANCH_ENDPOINT = "/v1/prompts/:alias/branches/:name"
    EVALUATE_ENDPOINT = "/v1/evaluate"

    EVALUATE_THREAD_ENDPOINT = "/v1/evaluate/threads/:threadId"
    EVALUATE_TRACE_ENDPOINT = "/v1/evaluate/traces/:traceUuid"
    EVALUATE_SPAN_ENDPOINT = "/v1/evaluate/spans/:spanUuid"

    METRICS_ENDPOINT = "/v1/metrics"
    METRIC_ENDPOINT = "/v1/metric/:name"

    GOVERNANCE_ASSESS_ENDPOINT = "/v1/governance/assess"


def _sanitize_body(obj):
    """Recursively replace non-finite floats (NaN, Inf, -Inf) with None.

    Python's json.dumps() happily serializes float('nan') as the
    literal token ``NaN`` which is **not** valid JSON and causes
    server-side parsing failures.  This helper walks any dict/list
    structure and neutralises those values before the payload is
    handed to the HTTP layer.
    """
    if isinstance(obj, float):
        return None if not math.isfinite(obj) else obj
    if isinstance(obj, dict):
        return {k: _sanitize_body(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)):
        return [_sanitize_body(v) for v in obj]
    return obj


class Api:
    def __init__(
        self,
        api_key: Optional[str] = None,
        verify_ssl: Optional[bool] = None,
    ):
        if api_key is None:
            api_key = get_confident_api_key()

        if not api_key:
            raise ValueError(
                f"No Confident API key found. Please run `deepeval login` or set the {CONFIDENT_API_KEY_ENV_VAR} environment variable in the CLI."
            )

        self.api_key = api_key
        self._headers = {
            "Content-Type": "application/json",
            "CONFIDENT-API-KEY": api_key,
            "X-DeepEval-Version": deepeval.__version__,
        }
        self.base_api_url = get_base_api_url()

        # Resolve SSL verification. Explicit constructor argument wins;
        # otherwise fall back to the CONFIDENT_DISABLE_SSL setting.
        if verify_ssl is None:
            verify_ssl = not bool(get_settings().CONFIDENT_DISABLE_SSL)
        self.verify_ssl = verify_ssl

    @retry(
        wait=wait_exponential_jitter(initial=1, exp_base=2, jitter=2, max=10),
        retry=retry_if_exception_type(retryable_exceptions),
        after=log_retry_error,
    )
    def _http_request(
        self, method: str, url: str, headers=None, json=None, params=None
    ):
        session = requests.Session()
        return session.request(
            method=method,
            url=url,
            headers=headers,
            json=json,
            params=params,
            verify=self.verify_ssl,
        )

    def _handle_response(
        self, response_data: Union[dict, Any]
    ) -> Tuple[Any, Optional[str]]:
        if not isinstance(response_data, dict):
            return response_data, None

        try:
            api_response = ApiResponse(**response_data)
        except Exception:
            return response_data, None

        if api_response.deprecated:
            deprecation_msg = "You are using a deprecated API endpoint. Please update your deepeval version."
            if api_response.link:
                deprecation_msg += f" See: {api_response.link}"
            logging.warning(deprecation_msg)

        if not api_response.success:
            error_message = api_response.error or "Request failed"
            raise ConfidentApiError(error_message, api_response.link)

        return api_response.data, api_response.link

    def send_request(
        self,
        method: HttpMethods,
        endpoint: Endpoints,
        body=None,
        params=None,
        url_params=None,
    ) -> Tuple[Any, Optional[str]]:
        url = f"{self.base_api_url}{endpoint.value}"

        # Replace URL parameters if provided
        if url_params:
            for key, value in url_params.items():
                placeholder = f":{key}"
                if placeholder in url:
                    url = url.replace(placeholder, str(value))

        if body is not None:
            body = _sanitize_body(body)

        res = self._http_request(
            method=method.value,
            url=url,
            headers=self._headers,
            json=body,
            params=params,
        )

        if res.status_code == 200:
            try:
                response_data = res.json()
                return self._handle_response(response_data)
            except ValueError:
                return res.text, None
        else:
            try:
                error_data = res.json()
                return self._handle_response(error_data)
            except (ValueError, ConfidentApiError) as e:
                if isinstance(e, ConfidentApiError):
                    raise e
                error_message = (
                    error_data.get("error", res.text)
                    if "error_data" in locals()
                    else res.text
                )
                raise Exception(error_message)

    async def a_send_request(
        self,
        method: HttpMethods,
        endpoint: Endpoints,
        body=None,
        params=None,
        url_params=None,
    ) -> Tuple[Any, Optional[str]]:
        url = f"{self.base_api_url}{endpoint.value}"

        if url_params:
            for key, value in url_params.items():
                placeholder = f":{key}"
                if placeholder in url:
                    url = url.replace(placeholder, str(value))

        if body is not None:
            body = _sanitize_body(body)

        async with aiohttp.ClientSession() as session:
            async with session.request(
                method=method.value,
                url=url,
                headers=self._headers,
                json=body,
                params=params,
                ssl=self.verify_ssl,
            ) as res:
                if res.status == 200:
                    try:
                        response_data = await res.json()
                        return self._handle_response(response_data)
                    except aiohttp.ContentTypeError:
                        return await res.text(), None
                else:
                    try:
                        error_data = await res.json()
                        return self._handle_response(error_data)
                    except (aiohttp.ContentTypeError, ConfidentApiError) as e:
                        if isinstance(e, ConfidentApiError):
                            raise e
                        error_message = (
                            error_data.get("error", await res.text())
                            if "error_data" in locals()
                            else await res.text()
                        )
                        raise Exception(error_message)