"""
HTTP + WebSocket client for security audit via api.aicippy.com + aivedha.ai.

This client communicates with the AiCippy security audit API which proxies
requests to aivedha.ai's AI security analysis engine. Authentication uses
the same auth.aivibe.cloud JWT token but does NOT require CHAKRA plan
validation -- security audit is a cross-service feature available to all
authenticated users.

Architecture:
    CLI -> SecurityAuditClient -> api.aicippy.com/api/v1/security/audit
                                        |
                                        v
                                   aivedha.ai AI engine
                                        |
                                        v
                              ws.aicippy.com (streamed results)
"""

from __future__ import annotations

import json
from typing import TYPE_CHECKING, Any

import httpx
import websockets
from tenacity import (
    retry,
    retry_if_exception_type,
    stop_after_attempt,
    wait_exponential,
)
from websockets.exceptions import ConnectionClosed, WebSocketException

from aicippy.security.models import (
    AuditFinding,
    AuditStatus,
    AuditSummary,
    SecurityAuditResult,
)
from aicippy.utils.logging import get_logger

if TYPE_CHECKING:
    from collections.abc import Callable

logger = get_logger(__name__)

_RETRY_KWARGS: dict[str, Any] = {
    "stop": stop_after_attempt(3),
    "wait": wait_exponential(multiplier=1, min=1, max=4),
    "retry": retry_if_exception_type((httpx.ConnectError, httpx.TimeoutException)),
    "reraise": True,
}


class SecurityAuditClient:
    """Client for security audit via api.aicippy.com + aivedha.ai AI analysis.

    Special auth: Uses the same auth.aivibe.cloud JWT token but does NOT
    require CHAKRA plan validation. This is a cross-service feature
    powered by aivedha.ai's AI security analysis engine.

    Args:
        auth_token: Bearer token from Cognito authentication.
        base_url: AiCippy API base URL (default: https://api.aicippy.com).
        ws_url: WebSocket URL for streaming results (default: wss://ws.aicippy.com).
        timeout: HTTP request timeout in seconds.
    """

    __slots__ = ("_base_url", "_client", "_ws_url")

    def __init__(
        self,
        auth_token: str,
        base_url: str = "https://api.aicippy.com",
        ws_url: str = "wss://ws.aicippy.com",
        timeout: float = 60.0,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._ws_url = ws_url.rstrip("/")
        self._client = httpx.AsyncClient(
            base_url=self._base_url,
            headers={
                "Authorization": f"Bearer {auth_token}",
                "Content-Type": "application/json",
                "Accept": "application/json",
                "User-Agent": "aicippy-cli/1.2.0",
            },
            timeout=httpx.Timeout(timeout),
        )

    # ── Start Audit ────────────────────────────────────────────────────

    @retry(**_RETRY_KWARGS)
    async def start_audit(
        self,
        scan_scope: str = "full",
        target_files: list[str] | None = None,
    ) -> str:
        """Start a security audit.

        POST /api/v1/security/audit

        The backend forwards to aivedha.ai for AI-based analysis.
        No plan validation check is performed (special case for security audit).

        Args:
            scan_scope: Scope of the audit (full, targeted, quick).
            target_files: Optional list of specific file paths to audit.

        Returns:
            audit_id: Unique identifier for tracking the audit.

        Raises:
            httpx.HTTPStatusError: On non-2xx response.
        """
        payload: dict[str, Any] = {"scan_scope": scan_scope}
        if target_files is not None:
            payload["target_files"] = target_files

        logger.info(
            "security_audit_start_request",
            scan_scope=scan_scope,
            target_file_count=len(target_files) if target_files else 0,
        )

        response = await self._client.post(
            "/api/v1/security/audit",
            json=payload,
        )
        response.raise_for_status()

        audit_id = response.json().get("audit_id", "")
        logger.info("security_audit_started", audit_id=audit_id)
        return audit_id

    # ── Get Audit Status ───────────────────────────────────────────────

    @retry(**_RETRY_KWARGS)
    async def get_audit_status(self, audit_id: str) -> SecurityAuditResult:
        """Get current audit status and any findings so far.

        GET /api/v1/security/audit/{audit_id}

        Args:
            audit_id: Identifier of the audit to query.

        Returns:
            SecurityAuditResult with current status and findings.

        Raises:
            httpx.HTTPStatusError: On non-2xx response.
        """
        response = await self._client.get(
            f"/api/v1/security/audit/{audit_id}",
        )
        response.raise_for_status()

        data = response.json()
        return SecurityAuditResult.from_dict(data)

    # ── Stream Audit Results via WebSocket ─────────────────────────────

    async def stream_audit_results(
        self,
        audit_id: str,
        callback: Callable[[AuditFinding], None],
    ) -> SecurityAuditResult:
        """Connect to WebSocket for real-time audit result streaming.

        Each message from the server is a JSON object representing either:
        - A finding (type=finding): parsed as AuditFinding and passed to callback
        - A status update (type=status): indicates scan progress
        - A completion signal (type=complete): includes full SecurityAuditResult

        Args:
            audit_id: Identifier of the audit to stream.
            callback: Callable invoked with each AuditFinding as it arrives.

        Returns:
            Final SecurityAuditResult after the stream completes.

        Raises:
            ConnectionError: If WebSocket connection fails.
            TimeoutError: If no completion signal is received within timeout.
        """
        ws_endpoint = f"{self._ws_url}/security/audit/{audit_id}"
        auth_header = self._client.headers.get("Authorization", "")

        logger.info(
            "security_audit_stream_connecting",
            audit_id=audit_id,
            ws_endpoint=ws_endpoint,
        )

        findings: list[AuditFinding] = []
        final_result: SecurityAuditResult | None = None

        try:
            async with websockets.connect(
                ws_endpoint,
                additional_headers={"Authorization": auth_header},
                ping_interval=30,
                ping_timeout=10,
                close_timeout=5,
            ) as ws:
                logger.info("security_audit_stream_connected", audit_id=audit_id)

                async for raw_message in ws:
                    if isinstance(raw_message, bytes):
                        raw_message = raw_message.decode("utf-8")

                    try:
                        data = json.loads(raw_message)
                    except json.JSONDecodeError as e:
                        logger.warning(
                            "security_audit_stream_invalid_json",
                            error=str(e),
                        )
                        continue

                    msg_type = data.get("type", "")

                    if msg_type == "finding":
                        finding = AuditFinding.from_dict(data.get("data", data))
                        findings.append(finding)
                        callback(finding)

                    elif msg_type == "status":
                        logger.info(
                            "security_audit_stream_status",
                            status=data.get("status", ""),
                            progress=data.get("progress", 0),
                        )

                    elif msg_type == "complete":
                        result_data = data.get("data", data)
                        # Merge streamed findings into the final result
                        if "findings" not in result_data:
                            result_data["findings"] = [f.to_dict() for f in findings]
                        final_result = SecurityAuditResult.from_dict(result_data)
                        logger.info(
                            "security_audit_stream_complete",
                            audit_id=audit_id,
                            total_findings=len(final_result.findings),
                        )
                        break

                    elif msg_type == "error":
                        error_msg = data.get("message", "Unknown audit error")
                        logger.error(
                            "security_audit_stream_error",
                            audit_id=audit_id,
                            error=error_msg,
                        )
                        final_result = SecurityAuditResult(
                            audit_id=audit_id,
                            status=AuditStatus.ERROR,
                            summary=None,
                            findings=findings,
                            recommendations=[],
                            scan_scope="unknown",
                            timestamp="",
                            ai_confidence=0.0,
                        )
                        break

        except (ConnectionClosed, WebSocketException) as e:
            logger.error(
                "security_audit_stream_connection_error",
                audit_id=audit_id,
                error=str(e),
            )
            # If we got partial findings, return them with error status
            if findings:
                return SecurityAuditResult(
                    audit_id=audit_id,
                    status=AuditStatus.ERROR,
                    summary=None,
                    findings=findings,
                    recommendations=["Audit stream was interrupted. Results may be partial."],
                    scan_scope="unknown",
                    timestamp="",
                    ai_confidence=0.0,
                )
            raise ConnectionError(f"WebSocket connection failed for audit {audit_id}: {e}") from e

        if final_result is not None:
            return final_result

        # Fallback: stream ended without explicit completion
        return SecurityAuditResult(
            audit_id=audit_id,
            status=AuditStatus.COMPLETE if findings else AuditStatus.ERROR,
            summary=AuditSummary(
                total_findings=len(findings),
                critical_count=sum(1 for f in findings if f.severity == "critical"),
                high_count=sum(1 for f in findings if f.severity == "high"),
                medium_count=sum(1 for f in findings if f.severity == "medium"),
                low_count=sum(1 for f in findings if f.severity == "low"),
                info_count=sum(1 for f in findings if f.severity == "info"),
                scan_duration_seconds=0.0,
                files_scanned=0,
                lines_scanned=0,
            )
            if findings
            else None,
            findings=findings,
            recommendations=[],
            scan_scope="full",
            timestamp="",
            ai_confidence=0.0,
        )

    # ── Get Fix Recommendation ─────────────────────────────────────────

    @retry(**_RETRY_KWARGS)
    async def get_fix_recommendation(self, finding_id: str) -> dict[str, Any]:
        """Get detailed AI-generated fix recommendation from aivedha.ai.

        GET /api/v1/security/audit/findings/{finding_id}/fix

        Args:
            finding_id: Identifier of the finding to get a fix for.

        Returns:
            Dictionary containing:
                - fix_code: Suggested code fix
                - explanation: Detailed explanation of the fix
                - references: Related documentation links

        Raises:
            httpx.HTTPStatusError: On non-2xx response.
        """
        response = await self._client.get(
            f"/api/v1/security/audit/findings/{finding_id}/fix",
        )
        response.raise_for_status()

        result: dict[str, Any] = response.json()
        logger.info(
            "security_audit_fix_retrieved",
            finding_id=finding_id,
        )
        return result

    # ── Lifecycle ──────────────────────────────────────────────────────

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

    async def __aenter__(self) -> SecurityAuditClient:
        return self

    async def __aexit__(self, *args: Any) -> None:
        await self.close()
