import os
import json
import time
import warnings
import google.generativeai as genai
from core import utils

# Suppress Vertex AI warnings
warnings.filterwarnings("ignore", category=UserWarning, module="vertexai")

try:
    import vertexai
    from vertexai.preview.vision_models import ImageGenerationModel as VertexImageModel
    HAS_VERTEX = True
except ImportError:
    HAS_VERTEX = False

try:
    from google.auth.transport.requests import Request
    from google.oauth2.credentials import Credentials
    from google_auth_oauthlib.flow import InstalledAppFlow
    HAS_OAUTH = True
except ImportError:
    HAS_OAUTH = False

model_logic = None
model_writer = None
model_artist = None
model_image = None
logic_model_name = "models/gemini-1.5-pro"
writer_model_name = "models/gemini-1.5-flash"
artist_model_name = "models/gemini-1.5-flash"
image_model_name = None
image_model_source = "None"


class ResilientModel:
    def __init__(self, name, safety_settings, role):
        self.name = name
        self.safety_settings = safety_settings
        self.role = role
        self.model = genai.GenerativeModel(name, safety_settings=safety_settings)

    def update(self, name):
        self.name = name
        self.model = genai.GenerativeModel(name, safety_settings=self.safety_settings)

    def generate_content(self, *args, **kwargs):
        retries = 0
        max_retries = 3
        base_delay = 5

        while True:
            try:
                return self.model.generate_content(*args, **kwargs)
            except Exception as e:
                err_str = str(e).lower()
                is_retryable = "429" in err_str or "quota" in err_str or "500" in err_str or "503" in err_str or "504" in err_str or "deadline" in err_str or "internal error" in err_str
                if is_retryable and retries < max_retries:
                    delay = base_delay * (2 ** retries)
                    utils.log("SYSTEM", f"⚠️ Quota error on {self.role} ({self.name}). Retrying in {delay}s...")
                    time.sleep(delay)

                    if retries == 0:
                        utils.log("SYSTEM", "Attempting to re-optimize models to find alternative...")
                        from ai import setup as _setup
                        _setup.init_models(force=True)

                    retries += 1
                    continue
                raise e