from .base import BaseProvider
import briefcase_ai as bf
import os
from typing import Any, Dict, Optional

class GeminiProvider(BaseProvider):
    """
    Google Gemini Provider supporting both AI Studio and Vertex AI.
    Models: gemini-1.5-pro, gemini-1.5-flash, gemini-2.0-flash, gemini-3-experimental
    """
    def __init__(self, model_id: str, **kwargs):
        super().__init__(model_id, **kwargs)
        self.use_vertex = kwargs.get("use_vertex", False)
        self._client = self._init_sdk()

    def _init_client(self):
        # High level init handled by _init_sdk
        pass

    def _init_sdk(self):
        if self.use_vertex:
            import vertexai
            from vertexai.generative_models import GenerativeModel
            vertexai.init(
                project=os.getenv("GOOGLE_CLOUD_PROJECT"),
                location=os.getenv("GOOGLE_CLOUD_LOCATION", "us-central1")
            )
            return GenerativeModel(self.model_id)
        else:
            import google.generativeai as genai
            genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
            return genai.GenerativeModel(self.model_id)

    def generate(self, prompt: str, **kwargs) -> str:
        # 1. Prepare Decision
        decision = bf.Decision("gemini_generation")
        decision.add_input(bf.Input("prompt", prompt, "string"))
        
        params = bf.ModelParameters(self.model_id)
        params.with_provider("google-vertex" if self.use_vertex else "google-ai-studio")
        decision.with_model_parameters(params)

        try:
            # 2. Execute
            response = self._client.generate_content(prompt, **kwargs)
            
            # 3. Extract Telemetry & Safety
            usage = getattr(response, 'usage_metadata', None)
            if usage:
                decision.add_tag("prompt_tokens", str(usage.prompt_token_count))
                decision.add_tag("candidates_tokens", str(usage.candidates_token_count))
                decision.add_tag("total_tokens", str(usage.total_token_count))

            # 4. Handle Safety Filters
            finish_reason = "UNKNOWN"
            if hasattr(response, 'candidates') and response.candidates:
                finish_reason = str(response.candidates[0].finish_reason)
                if "SAFETY" in finish_reason:
                    decision.add_tag("audit_reason", "SAFETY_FILTER_TRIGGERED")
            
            response_text = response.text
            
            # 5. Capture Output
            output = bf.Output("completion", response_text, "string")
            output.with_confidence(1.0) # Gemini doesn't expose confidence directly
            decision.add_output(output)
            
            bf.save(decision)
            return response_text

        except Exception as e:
            decision.with_error(str(e), type(e).__name__)
            bf.save(decision)
            raise
