"""Microsoft Defender ATP API client."""
import logging
import time
from typing import Any, Generator

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

from .constants import API_BASE, API_TIMEOUT, MAX_RETRIES, SSL_RETRIES, TOKEN_REFRESH_BUFFER

logger = logging.getLogger(__name__)


class DefenderClient:
    """Client for Microsoft Defender ATP API."""

    def __init__(self, client_id: str, client_secret: str, tenant_id: str):
        if not all([client_id, client_secret, tenant_id]):
            raise ValueError("client_id, client_secret, and tenant_id are required")

        self._client_id = client_id
        self._client_secret = client_secret
        self._tenant_id = tenant_id
        self._token: str | None = None
        self._token_expires_at: float = 0
        self._session = self._create_session()
        self._refresh_token()

    def _create_session(self) -> requests.Session:
        """Create session with retry configuration."""
        session = requests.Session()
        retry = Retry(
            total=MAX_RETRIES,
            backoff_factor=1,
            status_forcelist=[429, 500, 502, 503, 504],
        )
        session.mount("https://", HTTPAdapter(max_retries=retry))
        return session

    def _refresh_token(self) -> None:
        """Refresh OAuth token."""
        token_url = f"https://login.microsoftonline.com/{self._tenant_id}/oauth2/v2.0/token"
        response = self._session.post(
            token_url,
            data={
                "grant_type": "client_credentials",
                "client_id": self._client_id,
                "client_secret": self._client_secret,
                "scope": f"{API_BASE}/.default",
            },
            timeout=API_TIMEOUT,
        )
        response.raise_for_status()
        data = response.json()
        self._token = data["access_token"]
        expires_in = data.get("expires_in", 3600)
        self._token_expires_at = time.time() + expires_in - TOKEN_REFRESH_BUFFER
        logger.info(f"Authenticated (token expires in {expires_in}s)")

    def _headers(self) -> dict[str, str]:
        """Get headers with fresh token."""
        if time.time() >= self._token_expires_at:
            self._refresh_token()
        return {"Authorization": f"Bearer {self._token}", "Content-Type": "application/json"}

    def _request_with_retry(self, method: str, url: str, **kwargs: Any) -> requests.Response:
        """Make request with SSL/connection retry."""
        for attempt in range(1, SSL_RETRIES + 1):
            try:
                response = getattr(self._session, method)(
                    url, headers=self._headers(), timeout=API_TIMEOUT, **kwargs
                )
                response.raise_for_status()
                return response
            except (requests.exceptions.SSLError, requests.exceptions.ConnectionError) as e:
                if attempt >= SSL_RETRIES:
                    logger.error(f"Connection error after {SSL_RETRIES} retries: {e}")
                    raise
                sleep_time = 2**attempt
                logger.warning(f"Connection error (attempt {attempt}/{SSL_RETRIES}), retrying in {sleep_time}s")
                time.sleep(sleep_time)
        raise RuntimeError("Request failed after retries")

    def query(self, kql: str) -> Generator[dict[str, Any], None, None]:
        """Execute KQL query and yield results."""
        url = f"{API_BASE}/api/advancedqueries/run"
        response = self._request_with_retry("post", url, json={"Query": kql})
        data = response.json()
        yield from data.get("Results", [])
