"""
LLM helper using Vertex AI to bill against GCP credits.

Use this instead of the public Gemini API to ensure all usage
bills against your $25k GCP Ecosystem Partner credits.
"""

from google import genai
from typing import Optional
import os

# Default to your project with credits
DEFAULT_PROJECT = "aqueous-encoder-478015-r7"
DEFAULT_LOCATION = "us-central1"


class VertexLLM:
    """LLM client that bills against GCP Vertex AI credits."""

    def __init__(
        self,
        project_id: Optional[str] = None,
        location: Optional[str] = None,
        model: str = "gemini-2.0-flash-exp"
    ):
        """
        Initialize Vertex AI LLM client.

        Args:
            project_id: GCP project ID (defaults to aqueous-encoder-478015-r7)
            location: GCP region (defaults to us-central1)
            model: Model to use (defaults to gemini-2.0-flash-exp)
        """
        self.project_id = project_id or os.getenv("GCP_PROJECT_ID", DEFAULT_PROJECT)
        self.location = location or os.getenv("GCP_LOCATION", DEFAULT_LOCATION)
        self.model = model

        # Create client with Vertex AI enabled
        self.client = genai.Client(
            vertexai=True,
            project=self.project_id,
            location=self.location
        )

    def generate(
        self,
        prompt: str,
        temperature: float = 0.7,
        max_tokens: Optional[int] = None
    ) -> str:
        """
        Generate text using Gemini via Vertex AI.

        Args:
            prompt: The prompt to send to the model
            temperature: Sampling temperature (0.0 to 1.0)
            max_tokens: Maximum tokens to generate

        Returns:
            Generated text response
        """
        config = {}
        if temperature is not None:
            config['temperature'] = temperature
        if max_tokens is not None:
            config['max_output_tokens'] = max_tokens

        response = self.client.models.generate_content(
            model=self.model,
            contents=prompt,
            config=config if config else None
        )

        return response.text

    def generate_structured(
        self,
        prompt: str,
        system_instruction: Optional[str] = None,
        temperature: float = 0.7
    ) -> str:
        """
        Generate structured output with optional system instruction.

        Args:
            prompt: The prompt to send to the model
            system_instruction: Optional system instruction
            temperature: Sampling temperature

        Returns:
            Generated text response
        """
        # Combine system instruction with prompt if provided
        full_prompt = prompt
        if system_instruction:
            full_prompt = f"{system_instruction}\n\n{prompt}"

        return self.generate(
            prompt=full_prompt,
            temperature=temperature
        )


# Convenience function for quick usage
def ask_gemini(prompt: str, **kwargs) -> str:
    """
    Quick helper to ask Gemini a question via Vertex AI.

    Args:
        prompt: Question or prompt
        **kwargs: Additional arguments for VertexLLM.generate()

    Returns:
        Response text

    Example:
        >>> from fmatch.llm_vertex import ask_gemini
        >>> response = ask_gemini("What are the top 3 B2B enrichment providers?")
    """
    llm = VertexLLM()
    return llm.generate(prompt, **kwargs)


if __name__ == "__main__":
    # Quick test
    print("Testing Vertex AI LLM...")
    response = ask_gemini(
        "In 10 words or less, what makes a great spreadsheet tool?"
    )
    print(f"\nResponse: {response}")
    print("\n[This query was billed to your GCP credits]")
