"""
Microsoft Defender for Endpoint (TVM) API client — app-only.

Separate from Microsoft Graph: the threat & vulnerability management data
lives on api.securitycenter.microsoft.com with its own token resource
scope and app permission (Vulnerability.Read.All on WindowsDefenderATP).
Reuses the same Entra app (tenant/client/secret) as the Intune/Graph
integration; only the requested scope differs.

Returns REAL per-device CVEs (not pseudo) → they enrich like any CVE.
"""
import logging
import time
from datetime import datetime, timedelta
from typing import Any, Dict, List, Optional

import httpx
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type

logger = logging.getLogger(__name__)

MDE_BASE = "https://api.securitycenter.microsoft.com/api"
LOGIN_BASE = "https://login.microsoftonline.com"
MDE_SCOPE = "https://api.securitycenter.microsoft.com/.default"


class DefenderAPIError(Exception):
    pass


class DefenderAuthError(DefenderAPIError):
    pass


class DefenderClient:
    def __init__(self, tenant_id: str, client_id: str, client_secret: str, verify_ssl: bool = True):
        if not all([tenant_id, client_id, client_secret]):
            raise ValueError("Defender TVM requires tenant_id, client_id and client_secret.")
        self.tenant_id = tenant_id
        self.client_id = client_id
        self.client_secret = client_secret
        self._token: Optional[str] = None
        self._token_expires_at: Optional[datetime] = None
        self.client = httpx.Client(
            verify=verify_ssl,
            timeout=httpx.Timeout(30.0, connect=8.0),
            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5),
        )

    def _authenticate(self) -> str:
        url = f"{LOGIN_BASE}/{self.tenant_id}/oauth2/v2.0/token"
        data = {
            "grant_type": "client_credentials",
            "client_id": self.client_id,
            "client_secret": self.client_secret,
            "scope": MDE_SCOPE,
        }
        try:
            r = self.client.post(url, data=data)
            if r.status_code in (400, 401):
                detail = ""
                try:
                    detail = r.json().get("error_description", "")[:200]
                except Exception:
                    detail = r.text[:200]
                raise DefenderAuthError(f"token rejected ({r.status_code}): {detail}")
            r.raise_for_status()
            payload = r.json()
        except httpx.HTTPError as e:
            raise DefenderAuthError(f"token request failed: {e}") from e
        token = payload.get("access_token")
        if not token:
            raise DefenderAuthError("no access_token in token response")
        self._token = token
        self._token_expires_at = datetime.now() + timedelta(seconds=max(60, int(payload.get("expires_in", 3600)) - 120))
        return token

    def _ensure_token(self) -> str:
        if not self._token or not self._token_expires_at or datetime.now() >= self._token_expires_at:
            return self._authenticate()
        return self._token

    @retry(stop=stop_after_attempt(3), wait=wait_exponential(multiplier=1, min=1, max=8),
           retry=retry_if_exception_type((httpx.ReadTimeout, httpx.WriteTimeout)))
    def _get(self, url: str, params: Optional[Dict[str, Any]] = None,
             _tries: int = 4) -> Dict[str, Any]:
        # MDE allows ~100 calls/minute, and the per-machine vulnerability walk
        # is one call per machine — so a sync of any size hits 429, more so when
        # an app scan runs alongside it (tester: a wall of "429 Too Many
        # Requests"). Treating that as a hard error threw away everything the
        # sync had left to do; 429 is a "come back shortly", not a failure.
        # Microsoft states the wait in Retry-After, so honour it.
        token = self._ensure_token()
        for attempt in range(_tries):
            r = self.client.get(url, headers={"Authorization": f"Bearer {token}",
                                              "Accept": "application/json"},
                                params=params)
            if r.status_code != 429:
                break
            if attempt == _tries - 1:
                raise DefenderAPIError("Defender rate limit (429) — gave up after "
                                       f"{_tries} attempts")
            try:
                wait = float(r.headers.get("Retry-After", ""))
            except ValueError:
                wait = 0.0
            # No/!unusable header → back off geometrically instead of hammering.
            wait = min(max(wait, 2.0 * (2 ** attempt)), 60.0)
            logger.info("Defender 429 — retrying in %.0fs (attempt %d/%d)",
                        wait, attempt + 1, _tries)
            time.sleep(wait)
        if r.status_code >= 400:
            raise DefenderAPIError(f"Defender GET {url} -> {r.status_code}: {r.text[:200]}")
        return r.json()

    def _get_all(self, path: str, max_pages: int = 200) -> List[dict]:
        url = f"{MDE_BASE}{path}"
        out: List[dict] = []
        pages = 0
        while url and pages < max_pages:
            data = self._get(url)
            out.extend(data.get("value", []) or [])
            url = data.get("@odata.nextLink")
            pages += 1
        return out

    def get_machines(self) -> List[dict]:
        return self._get_all("/machines")

    def get_machine_vulnerabilities(self, machine_id: str) -> List[dict]:
        return self._get_all(f"/machines/{machine_id}/vulnerabilities")

    def get_software_vulnerabilities_by_machine(self) -> List[dict]:
        """Tenant-wide (device, CVE, software) assessment — the JSON-response
        export. One paginated call maps every machine+CVE to its affected
        software (the per-machine /vulnerabilities endpoint omits software).
        Rows carry deviceId/cveId + softwareVendor/softwareName/softwareVersion.
        Best-effort: returns [] if the tenant/plan doesn't expose it."""
        try:
            return self._get_all("/machines/SoftwareVulnerabilitiesByMachine")
        except DefenderAPIError as e:
            logger.warning("Defender software-vuln export unavailable: %s", e)
            return []

    def test_connection(self) -> dict:
        try:
            self._ensure_token()
        except DefenderAuthError as e:
            return {"ok": False, "step": "auth", "error": str(e)}
        try:
            data = self._get(f"{MDE_BASE}/machines", params={"$top": 1})
            return {"ok": True, "machine_sample": len(data.get("value", []) or [])}
        except DefenderAPIError as e:
            return {"ok": False, "step": "machines", "error": str(e)}

    def close(self) -> None:
        try:
            self.client.close()
        except Exception:
            pass
