# Copyright (c) 2026 MergeGuide, Inc. All rights reserved.
# SPDX-License-Identifier: Proprietary
"""
Drata GRC Platform Adapter

Integrates with Drata's API for compliance automation.
API Documentation: https://developers.drata.com/api-docs/
"""

from datetime import datetime, timedelta
from typing import Any

import httpx

from .base import (
    AuthToken,
    EvidenceType,
    GRCAdapter,
    GRCConfig,
    GRCEvidence,
    GRCPlatform,
    safe_raise_for_status,
    validate_resource_id,
)


class DrataAdapter(GRCAdapter):
    """
    Drata platform adapter.

    Supports:
    - OAuth 2.0 authentication
    - Evidence upload to controls
    - Test result submission
    - Control listing and status updates

    Rate Limit: 500 requests/minute
    """

    def __init__(self, config: GRCConfig):
        super().__init__(config)
        self._client = httpx.AsyncClient(timeout=30.0)

    @property
    def platform(self) -> GRCPlatform:
        return GRCPlatform.DRATA

    @property
    def base_url(self) -> str:
        return "https://api.drata.com"

    async def authenticate(self) -> AuthToken:
        """
        Authenticate with Drata using OAuth 2.0 client credentials.
        """
        response = await self._request_with_retry(
            self._client,
            "POST",
            f"{self.base_url}/oauth/token",
            data={
                "grant_type": "client_credentials",
                "client_id": self.config.client_id,
                "client_secret": self.config.client_secret,
            },
        )
        safe_raise_for_status(response, "Drata authentication")
        data = response.json()

        expires_at = datetime.utcnow() + timedelta(seconds=data.get("expires_in", 3600))

        return AuthToken(
            access_token=data["access_token"],
            token_type=data.get("token_type", "Bearer"),
            expires_in=data.get("expires_in", 3600),
            refresh_token=data.get("refresh_token"),
            scope=data.get("scope"),
            expires_at=expires_at,
        )

    async def refresh_token(self) -> AuthToken:
        """
        Refresh Drata authentication token.
        """
        if not self.config.refresh_token:
            return await self.authenticate()

        response = await self._request_with_retry(
            self._client,
            "POST",
            f"{self.base_url}/oauth/token",
            data={
                "grant_type": "refresh_token",
                "client_id": self.config.client_id,
                "client_secret": self.config.client_secret,
                "refresh_token": self.config.refresh_token,
            },
        )
        safe_raise_for_status(response, "Drata token refresh")
        data = response.json()

        expires_at = datetime.utcnow() + timedelta(seconds=data.get("expires_in", 3600))

        return AuthToken(
            access_token=data["access_token"],
            token_type=data.get("token_type", "Bearer"),
            expires_in=data.get("expires_in", 3600),
            refresh_token=data.get("refresh_token", self.config.refresh_token),
            expires_at=expires_at,
        )

    async def _get_headers(self) -> dict[str, str]:
        """Get headers with authentication token."""
        token = await self.ensure_authenticated()
        return {
            "Authorization": f"Bearer {token}",
            "Content-Type": "application/json",
        }

    async def upload_evidence(self, evidence: GRCEvidence, control_id: str) -> dict[str, Any]:
        """
        Upload evidence to a Drata control.

        Drata supports file uploads, URLs, and text evidence.
        """
        validate_resource_id(control_id, "control_id")
        headers = await self._get_headers()

        # Transform evidence to Drata format
        payload = self._transform_to_drata_evidence(evidence, control_id)

        response = await self._request_with_retry(
            self._client,
            "POST",
            f"{self.base_url}/controls/{control_id}/evidence",
            headers=headers,
            json=payload,
        )
        response.raise_for_status()

        return {
            "success": True,
            "evidence_id": response.json().get("id"),
            "control_id": control_id,
            "platform": self.platform.value,
        }

    async def submit_test_result(
        self,
        test_id: str,
        passed: bool,
        details: dict | None = None,
    ) -> dict[str, Any]:
        """
        Submit a test result to Drata.

        Test results directly affect control compliance status.
        """
        validate_resource_id(test_id, "test_id")
        headers = await self._get_headers()

        payload = {
            "testId": test_id,
            "passed": passed,
            "executedAt": datetime.utcnow().isoformat() + "Z",
            "details": details or {},
        }

        response = await self._request_with_retry(
            self._client,
            "POST",
            f"{self.base_url}/tests/{test_id}/results",
            headers=headers,
            json=payload,
        )
        response.raise_for_status()

        return {
            "success": True,
            "test_id": test_id,
            "result_id": response.json().get("id"),
            "passed": passed,
        }

    async def list_controls(self, framework: str | None = None) -> list[dict]:
        """
        List all controls from Drata.

        Args:
            framework: Optional framework filter (e.g., "SOC 2", "ISO 27001")
        """
        headers = await self._get_headers()

        params = {}
        if framework:
            params["framework"] = framework

        response = await self._request_with_retry(
            self._client,
            "GET",
            f"{self.base_url}/controls",
            headers=headers,
            params=params,
        )
        response.raise_for_status()

        controls = response.json().get("data", [])

        return [
            {
                "id": c.get("id"),
                "name": c.get("name"),
                "description": c.get("description"),
                "framework": c.get("framework"),
                "framework_control": c.get("referenceId"),
                "status": c.get("status"),
                "category": c.get("category"),
            }
            for c in controls
        ]

    async def get_control(self, control_id: str) -> dict:
        """
        Get details for a specific Drata control.
        """
        validate_resource_id(control_id, "control_id")
        headers = await self._get_headers()

        response = await self._request_with_retry(
            self._client,
            "GET",
            f"{self.base_url}/controls/{control_id}",
            headers=headers,
        )
        response.raise_for_status()

        c = response.json()
        return {
            "id": c.get("id"),
            "name": c.get("name"),
            "description": c.get("description"),
            "framework": c.get("framework"),
            "framework_control": c.get("referenceId"),
            "status": c.get("status"),
            "category": c.get("category"),
            "tests": c.get("tests", []),
            "evidence_count": len(c.get("evidence", [])),
        }

    async def update_control_status(
        self,
        control_id: str,
        status: str,
        evidence_ids: list[str] | None = None,
    ) -> dict:
        """
        Update a control's status in Drata.

        Note: In Drata, control status is typically determined by test results,
        not directly set. This method links evidence to controls.
        """
        validate_resource_id(control_id, "control_id")
        headers = await self._get_headers()

        # Link evidence to control if provided
        if evidence_ids:
            for evidence_id in evidence_ids:
                validate_resource_id(evidence_id, "evidence_id")
                await self._request_with_retry(
                    self._client,
                    "POST",
                    f"{self.base_url}/controls/{control_id}/evidence/{evidence_id}/link",
                    headers=headers,
                )

        return {
            "success": True,
            "control_id": control_id,
            "status": status,
            "linked_evidence": evidence_ids or [],
        }

    async def list_tests(self, control_id: str | None = None) -> list[dict]:
        """
        List automated tests in Drata.

        Args:
            control_id: Optional filter by control
        """
        headers = await self._get_headers()

        params = {}
        if control_id:
            params["controlId"] = control_id

        response = await self._request_with_retry(
            self._client,
            "GET",
            f"{self.base_url}/tests",
            headers=headers,
            params=params,
        )
        response.raise_for_status()

        tests = response.json().get("data", [])

        return [
            {
                "id": t.get("id"),
                "name": t.get("name"),
                "description": t.get("description"),
                "control_id": t.get("controlId"),
                "type": t.get("type"),
                "last_result": t.get("lastResult"),
            }
            for t in tests
        ]

    def _transform_to_drata_evidence(self, evidence: GRCEvidence, control_id: str) -> dict:
        """
        Transform MergeGuide evidence to Drata's evidence format.
        """
        # Determine evidence type mapping
        if evidence.evidence_type == EvidenceType.FILE:
            evidence_type = "FILE"
        elif evidence.evidence_type == EvidenceType.URL:
            evidence_type = "URL"
        else:
            evidence_type = "TEXT"

        payload = {
            "controlId": control_id,
            "evidenceType": evidence_type,
            "collectedAt": evidence.collected_at.isoformat() + "Z",
            "metadata": {
                "source": "MergeGuide",
                "evaluationId": evidence.evaluation_id,
                "repository": evidence.repository,
                "policyId": evidence.policy_id,
                "policyName": evidence.policy_name,
                "policyVersion": evidence.metadata.get("policy_version", "1.0.0"),
            },
        }

        if evidence_type == "FILE":
            payload["fileName"] = (
                evidence.file_name or f"mergeguide-evidence-{evidence.evidence_id}.json"
            )
            payload["fileContent"] = evidence.content  # Already base64 encoded
        elif evidence_type == "URL":
            payload["url"] = evidence.content
        else:
            payload["text"] = evidence.content

        return payload

    async def create_custom_connection(
        self,
        name: str,
        description: str,
    ) -> dict:
        """
        Create a custom connection in Drata for MergeGuide.

        Custom connections allow automated evidence collection.
        """
        headers = await self._get_headers()

        payload = {
            "name": name,
            "description": description,
            "type": "CUSTOM",
            "metadata": {
                "source": "MergeGuide",
                "version": "0.4.0",
            },
        }

        response = await self._request_with_retry(
            self._client,
            "POST",
            f"{self.base_url}/custom-connections",
            headers=headers,
            json=payload,
        )
        response.raise_for_status()

        return response.json()

    async def close(self):
        """Close the HTTP client."""
        await self._client.aclose()
