import logging
import time
from typing import Any, Optional
from urllib.parse import urlencode

import requests
from requests.adapters import HTTPAdapter
from requests.auth import HTTPBasicAuth
from urllib3.util.retry import Retry

logger = logging.getLogger(__name__)


class AuthenticationError(Exception):
    """Raised when API authentication fails."""


class PiDataMetrics:
    def __init__(
        self,
        client_id,
        client_secret,
        account_id,
        authorization_code: Optional[str] = None,
        redirect_uri: Optional[str] = None,
        refresh_token: Optional[str] = None,
        token_refresh_interval_seconds: int = 3600,
    ):
        self.client_id = client_id
        self.client_secret = client_secret
        self.account_id = account_id
        self.authorization_code = authorization_code
        self.redirect_uri = redirect_uri
        self.refresh_token = refresh_token
        self.token_refresh_interval_seconds = token_refresh_interval_seconds

        self.auth_url = "https://app.pi-datametrics.com/api/auth"
        self.base_url = f"https://app.pi-datametrics.com/api/accounts/{account_id}"
        self.timeout = (5, 60)
        self.session = self._build_session()

        self.access_token: Optional[str] = None
        self.token_obtained_at: float = 0.0

        self._authenticate_initial()

    @staticmethod
    def _build_session() -> requests.Session:
        session = requests.Session()
        retry = Retry(
            total=3,
            backoff_factor=0.5,
            status_forcelist=[429, 500, 502, 503, 504],
            allowed_methods=["GET", "POST"],
            raise_on_status=False,
        )
        adapter = HTTPAdapter(max_retries=retry)
        session.mount("http://", adapter)
        session.mount("https://", adapter)
        return session

    @property
    def headers(self) -> dict[str, str]:
        return {
            "Authorization": f"Bearer {self.access_token}",
            "Content-Type": "application/json",
        }

    def _set_token(self, payload: dict[str, Any]):
        token = payload.get("access_token")
        if not token:
            raise AuthenticationError("Authentication succeeded but no access_token was returned.")
        self.access_token = token
        self.refresh_token = payload.get("refresh_token", self.refresh_token)
        self.token_obtained_at = time.time()

    def _post_auth(self, data: dict[str, Any], use_basic_auth: bool = False):
        auth = HTTPBasicAuth(self.client_id, self.client_secret) if use_basic_auth else None
        response = self.session.post(
            self.auth_url,
            data=data,
            auth=auth,
            timeout=self.timeout,
        )
        response.raise_for_status()
        return response.json()

    def _authenticate_initial(self):
        if self.authorization_code and self.redirect_uri:
            self._authenticate_with_authorization_code()
            return
        self._authenticate_with_client_credentials()

    def _authenticate_with_authorization_code(self):
        data = {
            "grant_type": "authorization_code",
            "code": self.authorization_code,
            "redirect_uri": self.redirect_uri,
            "client_id": self.client_id,
            "client_secret": self.client_secret,
        }
        try:
            payload = self._post_auth(data, use_basic_auth=False)
            self._set_token(payload)
            logger.info("Authenticated with authorization_code grant.")
        except requests.exceptions.RequestException as exc:
            logger.exception("Authorization-code authentication failed")
            raise AuthenticationError(f"Authentication Failed: {exc}") from exc

    def _authenticate_with_client_credentials(self):
        data = {"grant_type": "client_credentials"}
        try:
            payload = self._post_auth(data, use_basic_auth=True)
            self._set_token(payload)
            logger.info("Authenticated with client_credentials grant.")
        except requests.exceptions.RequestException as exc:
            logger.exception("Client-credentials authentication failed")
            raise AuthenticationError(f"Authentication Failed: {exc}") from exc

    def _refresh_access_token(self):
        if not self.refresh_token:
            self._authenticate_with_client_credentials()
            return

        data = {
            "grant_type": "refresh_token",
            "refresh_token": self.refresh_token,
            "client_id": self.client_id,
            "client_secret": self.client_secret,
        }
        try:
            payload = self._post_auth(data, use_basic_auth=False)
            self._set_token(payload)
            logger.info("Access token refreshed using refresh_token grant.")
        except requests.exceptions.RequestException as exc:
            logger.warning("Refresh-token authentication failed, falling back to client_credentials: %s", exc)
            self._authenticate_with_client_credentials()

    def _ensure_token_valid(self, force_refresh: bool = False):
        if not self.access_token:
            self._authenticate_initial()
            return

        token_age = time.time() - self.token_obtained_at
        if force_refresh or token_age >= self.token_refresh_interval_seconds:
            self._refresh_access_token()

    def _request_json(self, url: str, params: Optional[dict[str, Any]] = None):
        self._ensure_token_valid()

        request_url = url
        if params:
            query_string = urlencode(params, doseq=True, safe="[]")
            request_url = f"{url}?{query_string}"
            params = None

        response = self.session.get(request_url, headers=self.headers, params=params, timeout=self.timeout)

        if response.status_code == 401:
            logger.warning("401 received, refreshing token and retrying request once.")
            self._ensure_token_valid(force_refresh=True)
            response = self.session.get(request_url, headers=self.headers, params=params, timeout=self.timeout)

        try:
            response.raise_for_status()
        except requests.exceptions.HTTPError as exc:
            logger.error("API Error: %s", exc)
            logger.error("Requested URL: %s", response.url)
            logger.error("Response Body: %s", response.text)
            raise

        return response.json().get("data", [])

    def fetch_endpoint(self, endpoint_path: str, params: Optional[dict[str, Any]] = None):
        url = f"{self.base_url}/{endpoint_path}"
        return self._request_json(url, params=params)

    def get_workspaces(self):
        return self.fetch_endpoint("workspaces")

    def get_stgs(self, workspace_id):
        return self.fetch_endpoint(f"workspaces/{workspace_id}/search-term-groups")

    def get_search_terms(self, workspace_id, stg_id):
        return self.fetch_endpoint(f"workspaces/{workspace_id}/search-term-groups/{stg_id}/search-terms")

    def get_bulk_serp_data(self, workspace_id, search_engine_id, period, **kwargs):
        params = {"search-engine-id": search_engine_id, "period": period}
        params.update(kwargs)
        return self.fetch_endpoint(f"workspaces/{workspace_id}/search-data/bulk-search-results", params=params)

    def get_bulk_volume(self, workspace_id, start_date=None, end_date=None):
        params = {}
        if start_date and end_date:
            params = {"start-period": start_date, "end-period": end_date}
        return self.fetch_endpoint(f"workspaces/{workspace_id}/volume-data/bulk-search-volume", params=params)

    def get_llm_mentions(self, workspace_id, search_engine_id, start_period, end_period, stg_ids=None):
        url = "https://app.pi-datametrics.com/api/data/llm/mentions"

        params = {
            "account-id": self.account_id,
            "workspace-id": workspace_id,
            "start-period": start_period,
            "end-period": end_period,
            "search-engine-id": search_engine_id,
        }

        if stg_ids:
            params["search-term-group-id[]"] = stg_ids

        return self._request_json(url, params=params)
