"""Vertex AI adapter for database analysis recommendations."""

import logging
import os
from typing import Optional

logger = logging.getLogger(__name__)


class VertexAIAdapter:
    """Adapter for Vertex AI Gemini integration."""

    def __init__(
        self,
        project: Optional[str] = None,
        region: Optional[str] = None,
        model: str = "gemini-1.5-pro",
        temperature: float = 0.3,
        max_tokens: int = 300,
        stream: bool = False,
    ) -> None:
        self.project = (
            project
            or os.getenv("VERTEX_PROJECT")
            or os.getenv("GCP_PROJECT")
            or os.getenv("GOOGLE_CLOUD_PROJECT")
        )
        self.region = (
            region
            or os.getenv("VERTEX_REGION")
            or os.getenv("GCP_REGION")
            or os.getenv("GOOGLE_CLOUD_REGION")
        )
        if not self.project or not self.region:
            raise ValueError(
                "Vertex AI requires project and region. Set VERTEX_PROJECT/GCP_PROJECT "
                "and VERTEX_REGION/GCP_REGION, or pass project/region explicitly."
            )

        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.stream = stream

        try:
            import vertexai
            from vertexai.generative_models import GenerationConfig, GenerativeModel

            self._vertexai = vertexai
            self._GenerationConfig = GenerationConfig
            self._GenerativeModel = GenerativeModel
        except ImportError as exc:  # pragma: no cover - optional dependency
            raise ImportError(
                "vertexai not installed. Install with: pip install "
                "google-cloud-aiplatform"
            ) from exc

        self._vertexai.init(project=self.project, location=self.region)
        self._client = self._GenerativeModel(self.model)
        self._generation_config = self._GenerationConfig(
            temperature=self.temperature,
            max_output_tokens=self.max_tokens,
        )

    def analyze_postgresql(
        self, prompt: str, system_instruction: Optional[str] = None
    ) -> str:
        return self._send_prompt(prompt, system_instruction)

    def analyze_mongodb(
        self, prompt: str, system_instruction: Optional[str] = None
    ) -> str:
        return self._send_prompt(prompt, system_instruction)

    def _send_prompt(
        self, prompt: str, system_instruction: Optional[str] = None
    ) -> str:
        try:
            if self.stream:
                return self._stream_response(prompt, system_instruction)
            else:
                response = self._client.generate_content(
                    prompt,
                    generation_config=self._generation_config,
                    system_instruction=system_instruction,
                )
                return response.text.strip() if hasattr(response, "text") else ""
        except Exception as exc:
            logger.error("Error calling Vertex AI: %s", exc)
            raise

    def _stream_response(
        self, prompt: str, system_instruction: Optional[str] = None
    ) -> str:
        """Stream response from Vertex AI and collect output

        Args:
            prompt: Prompt to send
            system_instruction: Optional system instruction

        Returns:
            Collected response text
        """
        full_response = []
        response = self._client.generate_content(
            prompt,
            generation_config=self._generation_config,
            system_instruction=system_instruction,
            stream=True,
        )

        for chunk in response:
            if hasattr(chunk, "text") and chunk.text:
                full_response.append(chunk.text)
                # Print chunks for live feedback
                print(chunk.text, end="", flush=True)

        print()  # New line after streaming
        return "".join(full_response)
