"""
Wazuh API Client

Secure integration with Wazuh Manager for vulnerability and agent management.

Features:
- JWT-based authentication
- Automatic token refresh
- Retry logic with exponential backoff
- Request throttling (Rate Limiting)
- SSL verification
- Comprehensive error handling

Wazuh API Documentation:
https://documentation.wazuh.com/current/user-manual/api/reference.html
"""
import os
import logging
import time
from typing import Optional, Dict, List, Any
from datetime import datetime, timedelta
import httpx
from tenacity import (
    retry,
    stop_after_attempt,
    wait_exponential,
    retry_if_exception_type
)

logger = logging.getLogger(__name__)


class WazuhAPIError(Exception):
    """Base exception for Wazuh API errors"""
    pass


class WazuhAuthenticationError(WazuhAPIError):
    """Authentication error"""
    pass


class WazuhRateLimitError(WazuhAPIError):
    """Rate limit exceeded"""
    pass


class WazuhClient:
    """
    Wazuh API Client with automatic authentication and retry logic

    Environment Variables:
        WAZUH_API_URL: Wazuh Manager URL (e.g. https://wazuh.example.com:55000)
        WAZUH_API_USERNAME: API Username
        WAZUH_API_PASSWORD: API Password
        WAZUH_VERIFY_SSL: Verify SSL certificate (default: true)
    """

    def __init__(
        self,
        base_url: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
        indexer_url: Optional[str] = None,
        indexer_username: Optional[str] = None,
        indexer_password: Optional[str] = None,
        verify_ssl: bool = True
    ):
        """
        Initializes the Wazuh API Client

        Args:
            base_url: Wazuh Manager URL (or from ENV: WAZUH_API_URL)
            username: API Username (or from ENV: WAZUH_API_USERNAME)
            password: API Password (or from ENV: WAZUH_API_PASSWORD)
            verify_ssl: Verify SSL certificate (or from ENV: WAZUH_VERIFY_SSL)
        """
        self.base_url = (base_url or os.getenv("WAZUH_API_URL", "")).rstrip("/")
        self.username = username or os.getenv("WAZUH_API_USERNAME")
        self.password = password or os.getenv("WAZUH_API_PASSWORD")
        self.verify_ssl = verify_ssl if verify_ssl is not None else os.getenv("WAZUH_VERIFY_SSL", "true").lower() == "true"

        # Indexer Connection (OpenSearch)
        self.indexer_url = (indexer_url or os.getenv("WAZUH_INDEXER_URL", "")).rstrip("/")
        self.indexer_username = indexer_username or os.getenv("WAZUH_INDEXER_USERNAME")
        self.indexer_password = indexer_password or os.getenv("WAZUH_INDEXER_PASSWORD")

        if not all([self.base_url, self.username, self.password]):
            raise ValueError(
                "Wazuh API Credentials missing. Please set WAZUH_API_URL, "
                "WAZUH_API_USERNAME and WAZUH_API_PASSWORD."
            )

        self._token: Optional[str] = None
        self._token_expires_at: Optional[datetime] = None

        # HTTP Client mit Timeouts und Connection-Pooling
        self.client = httpx.Client(
            base_url=self.base_url,
            verify=self.verify_ssl,
            timeout=httpx.Timeout(30.0, connect=10.0),
            limits=httpx.Limits(max_connections=10, max_keepalive_connections=5)
        )

        logger.info(f"Wazuh Client initialized: {self.base_url}")

    def _authenticate(self) -> str:
        """
        Authenticates against Wazuh API and retrieves JWT token

        Returns:
            JWT token as string

        Raises:
            WazuhAuthenticationError: If authentication fails
        """
        try:
            response = self.client.post(
                "/security/user/authenticate",
                auth=(self.username, self.password)
            )
            response.raise_for_status()

            data = response.json()
            token = data.get("data", {}).get("token")

            if not token:
                raise WazuhAuthenticationError("No token received in response")

            # Token Expiry (Wazuh: default 900 seconds = 15 minutes)
            # We refresh 2 minutes before
            self._token = token
            self._token_expires_at = datetime.now() + timedelta(minutes=13)

            logger.info("Wazuh authentication successful")
            return token

        except httpx.HTTPStatusError as e:
            logger.error(f"Wazuh Auth failed: {e.response.status_code} - {e.response.text}")
            raise WazuhAuthenticationError(f"Authentication failed: {e}")
        except Exception as e:
            logger.error(f"Wazuh Auth error: {e}")
            raise WazuhAuthenticationError(f"Authentication failed: {e}")

    def _ensure_authenticated(self) -> str:
        """
        Stellt sicher, dass ein gültiger Token vorhanden ist

        Returns:
            Gültiger JWT-Token
        """
        if not self._token or not self._token_expires_at or datetime.now() >= self._token_expires_at:
            logger.info("Token expired or missing - New Auth Request")
            return self._authenticate()

        return self._token

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=2, max=10),
        retry=retry_if_exception_type((httpx.TimeoutException, httpx.ConnectError))
    )
    def _request(
        self,
        method: str,
        endpoint: str,
        params: Optional[Dict[str, Any]] = None,
        json_data: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """
        Führt HTTP-Request mit Authentifizierung und Retry-Logik aus

        Args:
            method: HTTP-Methode (GET, POST, PUT, DELETE)
            endpoint: API-Endpoint (z.B. "/agents")
            params: URL-Parameter
            json_data: JSON-Body für POST/PUT

        Returns:
            JSON-Response als Dictionary

        Raises:
            WazuhAPIError: Bei API-Fehlern
            WazuhRateLimitError: Bei Rate-Limiting
        """
        token = self._ensure_authenticated()
        headers = {"Authorization": f"Bearer {token}"}

        try:
            response = self.client.request(
                method=method,
                url=endpoint,
                headers=headers,
                params=params,
                json=json_data
            )

            # Rate Limiting Check
            if response.status_code == 429:
                retry_after = int(response.headers.get("Retry-After", 60))
                logger.warning(f"Wazuh Rate Limit reached. Retry after {retry_after}s")
                raise WazuhRateLimitError(f"Rate Limit - Retry after {retry_after}s")

            response.raise_for_status()
            return response.json()

        except httpx.HTTPStatusError as e:
            logger.error(f"Wazuh API Error: {e.response.status_code} - {e.response.text}")
            raise WazuhAPIError(f"API Request failed: {e}")
        except Exception as e:
            logger.error(f"Wazuh Request Error: {e}")
            raise WazuhAPIError(f"Request failed: {e}")

    def _indexer_request(
        self,
        method: str,
        endpoint: str,
        json_data: Optional[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        """Führt Request gegen den Wazuh Indexer (OpenSearch) aus"""
        if not self.indexer_url:
            raise WazuhAPIError("Indexer URL not configured")
        
        try:
            response = self.client.request(
                method=method,
                url=f"{self.indexer_url}{endpoint}",
                auth=(self.indexer_username, self.indexer_password),
                json=json_data
            )
            response.raise_for_status()
            return response.json()
        except Exception as e:
            logger.error(f"Wazuh Indexer Request Error: {e}")
            raise WazuhAPIError(f"Indexer Request failed: {e}")

    # ============================================
    # Agent Management
    # ============================================

    def get_agents(
        self,
        status: Optional[str] = None,
        limit: int = 500,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """
        Holt Liste aller Wazuh Agents

        Args:
            status: Filter nach Status (active, disconnected, never_connected)
            limit: Maximale Anzahl Ergebnisse
            offset: Offset für Pagination

        Returns:
            Liste von Agent-Dictionaries
        """
        params = {"limit": limit, "offset": offset}
        if status:
            params["status"] = status

        response = self._request("GET", "/agents", params=params)
        return response.get("data", {}).get("affected_items", [])

    def get_agent_by_id(self, agent_id: str) -> Dict[str, Any]:
        """
        Holt Details zu einem spezifischen Agent

        Args:
            agent_id: Wazuh Agent ID (z.B. "001")

        Returns:
            Agent-Details als Dictionary
        """
        response = self._request("GET", f"/agents/{agent_id}")
        items = response.get("data", {}).get("affected_items", [])
        if not items:
            raise WazuhAPIError(f"Agent {agent_id} not found")
        return items[0]

    # ============================================
    # Vulnerability Management
    # ============================================

    def get_vulnerabilities(
        self,
        agent_id: str,
        limit: int = 5000,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """
        Holt alle Vulnerabilities eines Agents

        Args:
            agent_id: Wazuh Agent ID
            limit: Maximale Anzahl Ergebnisse
            offset: Offset für Pagination

        Returns:
            Liste von Vulnerability-Dictionaries

        Response-Format (Beispiel):
        {
            "cve": "CVE-2024-1234",
            "cvss": {
                "cvss3": {
                    "base_score": 9.8,
                    "vector": {
                        "attack_vector": "network",
                        "availability": "high",
                        ...
                    }
                }
            },
            "name": "openssh-server",
            "version": "1:8.2p1-4ubuntu0.5",
            "architecture": "amd64",
            "condition": "Package less than 1:8.2p1-4ubuntu0.11",
            "published": "2024-01-15T10:00:00Z",
            "severity": "Critical",
            "external_references": ["https://nvd.nist.gov/vuln/detail/CVE-2024-1234"]
        }
        """
        # Falls Indexer konfiguriert ist (Wazuh 4.8+), nutze diesen
        if self.indexer_url:
            return self.query_vulnerabilities_from_indexer(agent_id, limit, offset)

        params = {"limit": limit, "offset": offset}
        response = self._request("GET", f"/vulnerability/{agent_id}", params=params)
        return response.get("data", {}).get("affected_items", [])

    def query_vulnerabilities_from_indexer(
        self,
        agent_id: str,
        limit: int = 5000,
        offset: int = 0
    ) -> List[Dict[str, Any]]:
        """
        Holt Vulnerabilities aus dem Wazuh Indexer (OpenSearch)
        Erforderlich ab Wazuh 4.8.
        Filtert CVEs deren LETZTER Alert-Status "Solved" ist.

        Wichtig: Der alte Code hat alle CVEs gefiltert, die IRGENDWANN mal
        "Solved" waren. Das ist falsch, weil eine CVE erneut auftreten kann.
        Jetzt wird per top_hits nur der aktuellste Status pro CVE geprueft.
        """
        # Step 1: Get CVEs whose LATEST alert status is "Solved"
        solved_cves = set()
        try:
            solved_query = {
                "size": 0,
                "query": {
                    "bool": {
                        "must": [
                            {"term": {"agent.id": agent_id}},
                            {"exists": {"field": "data.vulnerability.status"}}
                        ]
                    }
                },
                "aggs": {
                    "per_cve": {
                        "terms": {"field": "data.vulnerability.cve", "size": 10000},
                        "aggs": {
                            "latest_status": {
                                "top_hits": {
                                    "size": 1,
                                    "sort": [{"@timestamp": {"order": "desc"}}],
                                    "_source": ["data.vulnerability.status"]
                                }
                            }
                        }
                    }
                }
            }
            solved_response = self._indexer_request(
                "POST",
                "/wazuh-alerts-*/_search",
                json_data=solved_query
            )
            buckets = solved_response.get("aggregations", {}).get("per_cve", {}).get("buckets", [])
            for bucket in buckets:
                latest_hit = bucket.get("latest_status", {}).get("hits", {}).get("hits", [])
                if latest_hit:
                    latest_status = latest_hit[0].get("_source", {}).get("data", {}).get("vulnerability", {}).get("status")
                    if latest_status == "Solved":
                        solved_cves.add(bucket["key"])
            if solved_cves:
                logger.info(f"Agent {agent_id}: Found {len(solved_cves)} CVEs with latest status 'Solved' to exclude")
        except Exception as e:
            logger.warning(f"Could not query solved CVEs from alerts: {e}")

        # Step 2: Query active vulnerabilities from states index.
        # Pagination: OpenSearch caps from+size at index.max_result_window
        # (default 10000), so we loop in pages of `limit` (default 5000)
        # until total_hits is drained or we hit MAX_TOTAL safety cap.
        # Prior behaviour was a single 5000-hit fetch — silently dropped
        # everything above 5000 on agents with very large finding sets
        # (tester saw 10000 hits returning only 5000). 50k cap covers
        # every real Wazuh fleet without unbounded memory growth.
        MAX_TOTAL = 50_000
        PAGE_SIZE = max(1, min(limit, 5000))
        hits: List[Dict[str, Any]] = []
        total_hits = 0
        page_offset = offset
        try:
            while True:
                query = {
                    "size": PAGE_SIZE,
                    "from": page_offset,
                    "query": {
                        "bool": {
                            "must": [{"term": {"agent.id": agent_id}}]
                        }
                    },
                    "track_total_hits": True,
                }
                response = self._indexer_request(
                    "POST",
                    "/wazuh-states-vulnerabilities-*/_search",
                    json_data=query,
                )
                page_hits = response.get("hits", {}).get("hits", []) or []
                total_hits = response.get("hits", {}).get("total", {}).get("value", 0)
                hits.extend(page_hits)
                logger.info(
                    f"Agent {agent_id}: page from={page_offset} returned "
                    f"{len(page_hits)} hits (total: {total_hits}, accumulated: {len(hits)})"
                )
                if not page_hits or len(hits) >= total_hits:
                    break
                if len(hits) >= MAX_TOTAL:
                    logger.warning(
                        f"Agent {agent_id}: hit MAX_TOTAL cap of {MAX_TOTAL}; "
                        f"{total_hits - len(hits)} vulnerabilities skipped"
                    )
                    break
                page_offset += PAGE_SIZE

            results = []
            skipped_no_cve = 0
            skipped_solved = 0
            for hit in hits:
                source = hit.get("_source", {})
                vuln = source.get("vulnerability", {})
                pkg = source.get("package", {})
                score = vuln.get("score", {})
                cve_id = vuln.get("id")

                if not cve_id:
                    skipped_no_cve += 1
                    continue

                # Skip if the LATEST alert status for this CVE is "Solved"
                if cve_id in solved_cves:
                    skipped_solved += 1
                    logger.debug(f"Agent {agent_id}: Skipping {cve_id} - latest alert status is Solved")
                    continue

                # Try multiple score field paths for compatibility
                raw_score = score.get("base") or score.get("base_score")
                # Wazuh-Indexer emits -1 / out-of-range placeholders for
                # packages it could not score. Clamp anything outside the
                # CVSSv3 spec range [0.0, 10.0] back to None so downstream
                # CPR maths cannot produce negative results (tester saw
                # CPR=-9.3 driven by cvss_score=-1).
                base_score = None
                try:
                    if raw_score is not None:
                        f = float(raw_score)
                        if 0.0 <= f <= 10.0:
                            base_score = f
                except (TypeError, ValueError):
                    base_score = None

                # Fixed-version extraction. Wazuh 4.x writes this in a
                # few different fields depending on the feed; try them
                # in order, then fall back to regex-parsing the
                # condition string ("Package less than X.Y.Z").
                fixed_version = (
                    vuln.get("fix")
                    or vuln.get("fixed_version")
                    or (pkg.get("fix") if isinstance(pkg, dict) else None)
                )
                if not fixed_version:
                    cond = vuln.get("condition") or (pkg.get("condition") if isinstance(pkg, dict) else None)
                    if isinstance(cond, str):
                        # "Package less than 1.2.3", "less than 4.5", etc.
                        import re as _re
                        m = _re.search(r"less than\s+([0-9][\w\.\-:+~]*)", cond, _re.IGNORECASE)
                        if m:
                            fixed_version = m.group(1)

                # Mapping auf das interne Format
                results.append({
                    "cve": cve_id,
                    "cvss": {
                        "cvss3": {
                            "base_score": base_score,
                        }
                    },
                    "name": pkg.get("name"),
                    "version": pkg.get("version"),
                    "fixed_version": fixed_version,
                    "severity": vuln.get("severity"),
                    "title": vuln.get("title") or vuln.get("description"),
                    "detected_at": vuln.get("detected_at") or source.get("@timestamp")
                })

            if skipped_no_cve:
                logger.warning(f"Agent {agent_id}: Skipped {skipped_no_cve} entries without CVE ID")
            if skipped_solved:
                logger.info(f"Agent {agent_id}: Skipped {skipped_solved} CVEs with latest status 'Solved'")

            logger.info(f"Agent {agent_id}: Returning {len(results)} vulnerabilities from indexer")
            return results
        except Exception as e:
            logger.error(f"Fehler beim Abfragen des Indexers fuer Agent {agent_id}: {e}")
            return []

    def get_vulnerability_summary(self, agent_id: str) -> Dict[str, int]:
        """
        Holt Vulnerability-Zusammenfassung für einen Agent

        Args:
            agent_id: Wazuh Agent ID

        Returns:
            Dictionary mit Counts pro Severity-Level
            Beispiel: {"critical": 5, "high": 12, "medium": 30, "low": 45}
        """
        response = self._request("GET", f"/vulnerability/{agent_id}/summary")
        return response.get("data", {}).get("affected_items", [{}])[0]

    # ============================================
    # Syscollector (System Inventory)
    # ============================================

    def trigger_syscollector_scan(self, agent_id: str) -> Dict[str, Any]:
        """
        Triggers a syscollector scan for an agent

        This is important for patch verification:
        After a system update, Wazuh must rescan installed packages
        to check if vulnerabilities have been fixed.

        Args:
            agent_id: Wazuh Agent ID

        Returns:
            Response with Scan Status

        Note:
            The scan runs asynchronously. It takes approx. 1-5 minutes until results are available.
            Status can be checked with get_agent_by_id() (lastKeepAlive timestamp).
        """
        # Wazuh provides no direct "Rescan" endpoint
        # Workaround: Restart the Agent (forces Syscollector-Rescan)
        # Check: PUT /agents/{agent_id}/restart
        response = self._request("PUT", f"/agents/{agent_id}/restart")
        logger.info(f"Syscollector scan triggered for agent {agent_id}")
        return response

    def get_packages(self, agent_id: str) -> List[Dict[str, Any]]:
        """
        Holt Liste aller installierten Packages eines Agents

        Args:
            agent_id: Wazuh Agent ID

        Returns:
            Liste von Package-Dictionaries
        """
        response = self._request("GET", f"/syscollector/{agent_id}/packages")
        return response.get("data", {}).get("affected_items", [])

    def get_os_info(self, agent_id: str) -> Dict[str, Any]:
        """
        Holt OS-Informationen eines Agents

        Args:
            agent_id: Wazuh Agent ID

        Returns:
            OS-Details (Name, Version, Kernel, etc.)
        """
        response = self._request("GET", f"/syscollector/{agent_id}/os")
        items = response.get("data", {}).get("affected_items", [])
        return items[0] if items else {}

    # ============================================
    # SCA — Security Configuration Assessment
    # ============================================

    def get_sca_policies(self, agent_id: str) -> List[Dict[str, Any]]:
        """
        List all SCA policies an agent has been evaluated against.

        Wazuh API: ``GET /sca/{agent_id}``
        Returns per-policy summary including pass/fail/invalid counts +
        score. Example item:
            {
              "policy_id": "cis_win11",
              "name": "CIS Microsoft Windows 11 Enterprise Benchmark v1.0.0",
              "description": "...",
              "pass": 142, "fail": 38, "invalid": 0, "total_checks": 180,
              "score": 78, "end_scan": "2026-05-19T01:00:00Z", ...
            }
        """
        response = self._request("GET", f"/sca/{agent_id}")
        return response.get("data", {}).get("affected_items", []) or []

    def get_sca_checks(
        self,
        agent_id: str,
        policy_id: str,
        limit: int = 500,
    ) -> List[Dict[str, Any]]:
        """
        Per-check detail for one (agent, policy).

        Wazuh API: ``GET /sca/{agent_id}/checks/{policy_id}``
        Used by the optional deep-dive endpoint — not called by the
        default compliance refresh which only persists per-policy
        summaries. Returns up to ``limit`` checks; the default 500 is
        enough for the largest CIS benchmarks (~250-400 checks).
        """
        response = self._request(
            "GET",
            f"/sca/{agent_id}/checks/{policy_id}",
            params={"limit": limit},
        )
        return response.get("data", {}).get("affected_items", []) or []

    # ============================================
    # Helper-Funktionen
    # ============================================

    def wait_for_scan_completion(
        self,
        agent_id: str,
        timeout: int = 600,
        check_interval: int = 10
    ) -> bool:
        """
        Wartet, bis ein Syscollector-Scan abgeschlossen ist

        Args:
            agent_id: Wazuh Agent ID
            timeout: Maximale Wartezeit in Sekunden (default: 10 Min)
            check_interval: Prüfintervall in Sekunden

        Returns:
            True wenn Scan abgeschlossen, False bei Timeout
        """
        start_time = time.time()
        initial_scan_time = self._get_last_scan_time(agent_id)

        while time.time() - start_time < timeout:
            time.sleep(check_interval)

            current_scan_time = self._get_last_scan_time(agent_id)
            if current_scan_time and current_scan_time > initial_scan_time:
                logger.info(f"Scan for agent {agent_id} completed")
                return True

        logger.warning(f"Scan timeout for agent {agent_id}")
        return False

    def _get_last_scan_time(self, agent_id: str) -> Optional[datetime]:
        """Holt Zeitpunkt des letzten Syscollector-Scans"""
        try:
            agent = self.get_agent_by_id(agent_id)
            last_keepalive = agent.get("lastKeepAlive")
            if last_keepalive:
                return datetime.fromisoformat(last_keepalive.replace("Z", "+00:00"))
        except Exception as e:
            logger.error(f"Error fetching scan time: {e}")
        return None

    def close(self):
        """Closes HTTP client and releases resources"""
        self.client.close()
        logger.info("Wazuh Client closed")

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# ============================================
# Convenience-Funktionen
# ============================================

def get_wazuh_client() -> WazuhClient:
    """
    Factory-Funktion für Dependency Injection in FastAPI

    Usage:
        @app.get("/sync")
        def sync_vulnerabilities(wazuh: WazuhClient = Depends(get_wazuh_client)):
            agents = wazuh.get_agents(status="active")
            ...
    """
    return WazuhClient()
