"""
OpenAI GPT LLM Client (supports both OpenAI and Azure OpenAI)

Based on C# OpenAIClient.cs
"""

import json
import time
from typing import Any

import httpx

from warden.shared.infrastructure.exceptions import ExternalServiceError
from warden.shared.infrastructure.resilience import resilient

from ..config import ProviderConfig
from ..registry import ProviderRegistry
from ..types import LlmProvider, LlmRequest, LlmResponse
from .base import ILlmClient


class OpenAIClient(ILlmClient):
    """OpenAI GPT client - supports both OpenAI and Azure OpenAI"""

    def __init__(self, config: ProviderConfig, provider: LlmProvider = LlmProvider.OPENAI):
        if not config.api_key:
            raise ValueError(f"{provider.value} API key is required")

        self._api_key = config.api_key
        self._provider = provider
        self._default_model = config.default_model or "gpt-4o"

        # Azure vs OpenAI endpoints
        if provider == LlmProvider.AZURE_OPENAI:
            if not config.endpoint:
                raise ValueError("Azure OpenAI endpoint is required")
            self._base_url = config.endpoint.rstrip("/")
            self._api_version = config.api_version or "2024-02-01"
        else:
            self._base_url = config.endpoint or "https://api.openai.com/v1"

        self._usage = {"total_tokens": 0, "prompt_tokens": 0, "completion_tokens": 0, "request_count": 0}

    @property
    def provider(self) -> LlmProvider:
        return self._provider

    def get_usage(self) -> dict[str, int]:
        """Get cumulative token usage."""
        return self._usage.copy()

    @resilient(name="openai_send", timeout_seconds=60.0, retry_max_attempts=3)
    async def send_async(self, request: LlmRequest) -> LlmResponse:
        start_time = time.time()

        try:
            from warden.llm.global_rate_limiter import GlobalRateLimiter

            limiter = await GlobalRateLimiter.get_instance()
            provider_name = "azure" if self._provider == LlmProvider.AZURE_OPENAI else "openai"
            await limiter.acquire(provider_name, tokens=request.max_tokens)

            headers = {"Content-Type": "application/json"}

            # Azure uses api-key header, OpenAI uses Authorization
            if self._provider == LlmProvider.AZURE_OPENAI:
                headers["api-key"] = self._api_key
                url = f"{self._base_url}/openai/deployments/{request.model or self._default_model}/chat/completions?api-version={self._api_version}"
            else:
                headers["Authorization"] = f"Bearer {self._api_key}"
                url = f"{self._base_url}/chat/completions"

            payload = {
                "messages": [
                    {"role": "system", "content": request.system_prompt},
                    {"role": "user", "content": request.user_message},
                ],
                "temperature": request.temperature,
                "max_tokens": request.max_tokens,
            }

            if self._provider != LlmProvider.AZURE_OPENAI:
                payload["model"] = request.model or self._default_model

            async with httpx.AsyncClient(timeout=request.timeout_seconds) as client:
                response = await client.post(url, headers=headers, json=payload)
                response.raise_for_status()
                result = response.json()

            duration_ms = int((time.time() - start_time) * 1000)

            if not result.get("choices"):
                return LlmResponse(
                    content="",
                    success=False,
                    error_message="No response from OpenAI",
                    provider=self.provider,
                    duration_ms=duration_ms,
                )

            usage = result.get("usage", {})
            prompt_tokens = usage.get("prompt_tokens", 0)
            completion_tokens = usage.get("completion_tokens", 0)
            total_tokens = usage.get("total_tokens", 0)

            # Update cumulative usage
            self._usage["prompt_tokens"] += prompt_tokens
            self._usage["completion_tokens"] += completion_tokens
            self._usage["total_tokens"] += total_tokens
            self._usage["request_count"] += 1

            return LlmResponse(
                content=result["choices"][0]["message"]["content"],
                success=True,
                provider=self.provider,
                model=result.get("model"),
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                total_tokens=total_tokens,
                duration_ms=duration_ms,
            )

        except (httpx.HTTPStatusError, httpx.RequestError) as e:
            duration_ms = int((time.time() - start_time) * 1000)
            return LlmResponse(
                content="",
                success=False,
                error_message=f"HTTP error: {e!s}",
                provider=self.provider,
                duration_ms=duration_ms,
            )
        except (json.JSONDecodeError, KeyError) as e:
            duration_ms = int((time.time() - start_time) * 1000)
            return LlmResponse(
                content="",
                success=False,
                error_message=f"JSON/Data error: {e!s}",
                provider=self.provider,
                duration_ms=duration_ms,
            )
        except Exception as e:
            # Last resort for truly unexpected errors
            duration_ms = int((time.time() - start_time) * 1000)
            return LlmResponse(
                content="",
                success=False,
                error_message=f"Unexpected error: {e!s}",
                provider=self.provider,
                duration_ms=duration_ms,
            )

    async def is_available_async(self) -> bool:
        """
        Check if provider is available.

        For Azure OpenAI, just check if credentials are configured.
        Making a test API call would waste tokens and time.
        """
        try:
            # Just verify we have the necessary credentials
            if self._provider == LlmProvider.AZURE_OPENAI:
                return bool(self._api_key and self._base_url)
            else:
                return bool(self._api_key)
        except (ValueError, AttributeError):
            return False

    async def complete_async(
        self, prompt: str, system_prompt: str = "You are a helpful coding assistant.", model: str | None = None
    ) -> LlmResponse:
        """
        Simple completion method for non-streaming requests.

        Args:
            prompt: User prompt
            system_prompt: System prompt (optional)
            model: Model name override (optional)

        Returns:
            LlmResponse object with content and token usage

        Raises:
            Exception: If request fails
        """
        request = LlmRequest(
            user_message=prompt,
            system_prompt=system_prompt,
            model=model or self._default_model,
            temperature=0.0,  # Idempotency
            max_tokens=2000,
            timeout_seconds=30.0,
        )

        response = await self.send_async(request)

        if not response.success:
            raise ExternalServiceError(f"LLM request failed: {response.error_message}")

        return response

    async def stream_completion_async(
        self, prompt: str, system_prompt: str = "You are a helpful coding assistant.", model: str | None = None
    ):
        """
        Streaming completion method using Server-Sent Events (SSE).

        Args:
            prompt: User prompt
            system_prompt: System prompt (optional)
            model: Model name override (optional)

        Yields:
            Completion chunks as they arrive from the API

        Note:
            Falls back to simulated streaming if true streaming fails
        """
        try:
            # Attempt true streaming
            async for chunk in self._stream_with_sse(prompt, system_prompt, model):
                yield chunk
        except Exception as e:
            # Fallback to simulated streaming on error
            from warden.shared.infrastructure.logging import get_logger

            logger = get_logger(__name__)
            logger.warning("openai_streaming_fallback", error=str(e), provider=self._provider.value)

            # Use non-streaming as fallback
            response = await self.complete_async(prompt, system_prompt, model)
            full_response = response.content

            # Simulate streaming by yielding in chunks
            chunk_size = 20
            for i in range(0, len(full_response), chunk_size):
                chunk = full_response[i : i + chunk_size]
                yield chunk

    async def _stream_with_sse(self, prompt: str, system_prompt: str, model: str | None = None):
        """
        Internal method for true SSE streaming.

        Args:
            prompt: User prompt
            system_prompt: System prompt
            model: Model name override (optional)

        Yields:
            Content chunks from the streaming response
        """
        headers = {"Content-Type": "application/json"}

        # Azure vs OpenAI authentication
        if self._provider == LlmProvider.AZURE_OPENAI:
            headers["api-key"] = self._api_key
            url = f"{self._base_url}/openai/deployments/{model or self._default_model}/chat/completions?api-version={self._api_version}"
        else:
            headers["Authorization"] = f"Bearer {self._api_key}"
            url = f"{self._base_url}/chat/completions"

        payload = {
            "messages": [{"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}],
            "temperature": 0.0,
            "max_tokens": 2000,
            "stream": True,  # Enable streaming
        }

        # Add model for non-Azure
        if self._provider != LlmProvider.AZURE_OPENAI:
            payload["model"] = model or self._default_model

        async with httpx.AsyncClient(timeout=120.0) as client:
            async with client.stream("POST", url, headers=headers, json=payload) as response:
                response.raise_for_status()

                # Process SSE stream
                async for line in response.aiter_lines():
                    line = line.strip()

                    # Skip empty lines and comments
                    if not line or line.startswith(":"):
                        continue

                    # SSE format: "data: {json}"
                    if line.startswith("data: "):
                        data_str = line[6:]  # Remove "data: " prefix

                        # Stream termination signal
                        if data_str == "[DONE]":
                            break

                        try:
                            chunk_data = json.loads(data_str)

                            # Extract content from delta
                            choices = chunk_data.get("choices", [])
                            if choices:
                                delta = choices[0].get("delta", {})
                                content = delta.get("content", "")
                                if content:
                                    yield content
                        except json.JSONDecodeError:
                            # Skip malformed JSON chunks
                            continue

    async def analyze_security_async(self, code_content: str, language: str) -> dict[str, Any]:
        """
        Analyze code for security vulnerabilities using LLM.

        Args:
            code_content: Source code to analyze
            language: Language of the code

        Returns:
            Dict containing findings list
        """
        from warden.shared.utils.json_parser import parse_json_from_llm
        from warden.shared.utils.token_utils import truncate_content_for_llm

        # Truncate code using token-aware truncation (not character-based)
        truncated_code = truncate_content_for_llm(
            code_content,
            max_tokens=1800,  # Reserve tokens for prompt and response
            preserve_start_lines=30,
            preserve_end_lines=15,
        )

        prompt = f"""
        You are a senior security researcher. Analyze this {language} code for critical vulnerabilities.
        Target vulnerabilities: SQL Injection, XSS, Hardcoded Secrets/Credentials, SSRF, CSRF, XXE, Insecure Deserialization, Path Traversal, and Command Injection.

        Ignore stylistic issues. Focus only on exploitable security flaws.

        Return a JSON object in this exact format:
        {{
            "findings": [
                {{
                    "severity": "critical|high|medium",
                    "message": "Short description",
                    "line_number": 1,
                    "detail": "Detailed explanation of the exploit vector",
                    "source": "Where tainted data enters, e.g. request.args['id'] (line 14)",
                    "sink": "Where tainted data is consumed unsafely, e.g. cursor.execute() (line 45)",
                    "data_flow": ["function_or_variable_names", "showing", "the path"]
                }}
            ]
        }}

        The source, sink, and data_flow fields are optional but highly valuable for triage.
        If no issues found, return {{ "findings": [] }}.

        Code:
        ```{language}
        {truncated_code}
        ```
        """

        try:
            response = await self.complete_async(
                prompt, system_prompt="You are a strict security auditor. Output valid JSON only."
            )
            if not response.success:
                return {"findings": []}

            parsed = parse_json_from_llm(response.content)
            if not parsed:
                return {"findings": []}

            # Enrich findings with MachineContext from structured LLM output
            self._enrich_findings_from_llm(parsed)
            return parsed
        except (RuntimeError, ValueError, json.JSONDecodeError) as e:
            # Fallback to empty findings on failure to prevent crash
            from warden.shared.infrastructure.logging import get_logger

            logger = get_logger(__name__)
            logger.error("llm_security_analysis_failed", error=str(e))
            return {"findings": []}


# Self-register with the registry
# OpenAI provider handles both OPENAI and AZURE_OPENAI with provider parameter
ProviderRegistry.register(LlmProvider.OPENAI, lambda config: OpenAIClient(config, LlmProvider.OPENAI))
ProviderRegistry.register(LlmProvider.AZURE_OPENAI, lambda config: OpenAIClient(config, LlmProvider.AZURE_OPENAI))
