app/vertex_ai_init.py

import json
import asyncio # Added for await
from google import genai
from credentials_manager import CredentialManager, parse_multiple_json_credentials
import config as app_config
from model_loader import refresh_models_config_cache # Import new model loader function

# VERTEX_EXPRESS_MODELS list is now dynamically loaded via model_loader
# The constant VERTEX_EXPRESS_MODELS previously defined here is removed.
# Consumers should use get_vertex_express_models() from model_loader.

# Global 'client' and 'get_vertex_client()' are removed.

async def init_vertex_ai(credential_manager_instance: CredentialManager) -> bool: # Made async
    """
    Initializes the credential manager with credentials from GOOGLE_CREDENTIALS_JSON (if provided)
    and verifies if any credentials (environment or file-based through the manager) are available.
    The CredentialManager itself handles loading file-based credentials upon its instantiation.
    This function primarily focuses on augmenting the manager with env var credentials.

    Returns True if any credentials seem available in the manager, False otherwise.
    """
    try:
        credentials_json_str = app_config.GOOGLE_CREDENTIALS_JSON_STR
        env_creds_loaded_into_manager = False

        if credentials_json_str:
            print("INFO: Found GOOGLE_CREDENTIALS_JSON environment variable. Attempting to load into CredentialManager.")
            try:
                # Attempt 1: Parse as multiple JSON objects
                json_objects = parse_multiple_json_credentials(credentials_json_str)
                if json_objects:
                    print(f"DEBUG: Parsed {len(json_objects)} potential credential objects from GOOGLE_CREDENTIALS_JSON.")
                    success_count = credential_manager_instance.load_credentials_from_json_list(json_objects)
                    if success_count > 0:
                        print(f"INFO: Successfully loaded {success_count} credentials from GOOGLE_CREDENTIALS_JSON into manager.")
                        env_creds_loaded_into_manager = True
                
                # Attempt 2: If multiple parsing/loading didn't add any, try parsing/loading as a single JSON object
                if not env_creds_loaded_into_manager:
                    print("DEBUG: Multi-JSON loading from GOOGLE_CREDENTIALS_JSON did not add to manager or was empty. Attempting single JSON load.")
                    try:
                        credentials_info = json.loads(credentials_json_str)
                        # Basic validation (CredentialManager's add_credential_from_json does more thorough validation)

                        if isinstance(credentials_info, dict) and \
                           all(field in credentials_info for field in ["type", "project_id", "private_key_id", "private_key", "client_email"]):
                            if credential_manager_instance.add_credential_from_json(credentials_info):
                                print("INFO: Successfully loaded single credential from GOOGLE_CREDENTIALS_JSON into manager.")
                                # env_creds_loaded_into_manager = True # Redundant, as this block is conditional on it being False
                            else:
                                print("WARNING: Single JSON from GOOGLE_CREDENTIALS_JSON failed to load into manager via add_credential_from_json.")
                        else:
                             print("WARNING: Single JSON from GOOGLE_CREDENTIALS_JSON is not a valid dict or missing required fields for basic check.")
                    except json.JSONDecodeError as single_json_err:
                        print(f"WARNING: GOOGLE_CREDENTIALS_JSON could not be parsed as a single JSON object: {single_json_err}.")
                    except Exception as single_load_err:
                        print(f"WARNING: Error trying to load single JSON from GOOGLE_CREDENTIALS_JSON into manager: {single_load_err}.")
            except Exception as e_json_env:
                # This catches errors from parse_multiple_json_credentials or load_credentials_from_json_list
                print(f"WARNING: Error processing GOOGLE_CREDENTIALS_JSON env var: {e_json_env}.")
        else:
            print("INFO: GOOGLE_CREDENTIALS_JSON environment variable not found.")

        # Attempt to pre-warm the model configuration cache
        print("INFO: Attempting to pre-warm model configuration cache during startup...")
        models_loaded_successfully = await refresh_models_config_cache()
        if models_loaded_successfully:
            print("INFO: Model configuration cache pre-warmed successfully.")
        else:
            print("WARNING: Failed to pre-warm model configuration cache during startup. It will be loaded lazily on first request.")
            # We don't necessarily fail the entire init_vertex_ai if model list fetching fails,
            # as credential validation might still be important, and model list can be fetched later.

        # CredentialManager's __init__ calls load_credentials_list() for files.
        # refresh_credentials_list() re-scans files and combines with in-memory (already includes env creds if loaded above).
        # The return value of refresh_credentials_list indicates if total > 0
        if credential_manager_instance.refresh_credentials_list():
            total_creds = credential_manager_instance.get_total_credentials()
            print(f"INFO: Credential Manager reports {total_creds} credential(s) available (from files and/or GOOGLE_CREDENTIALS_JSON).")
            
            # Optional: Attempt to validate one of the credentials by creating a temporary client.
            # This adds a check that at least one credential is functional.
            print("INFO: Attempting to validate a random credential by creating a temporary client...")
            temp_creds_val, temp_project_id_val = credential_manager_instance.get_random_credentials()
            if temp_creds_val and temp_project_id_val:
                try:
                    _ = genai.Client(vertexai=True, credentials=temp_creds_val, project=temp_project_id_val, location="global")
                    print(f"INFO: Successfully validated a credential from Credential Manager (Project: {temp_project_id_val}). Initialization check passed.")
                    return True
                except Exception as e_val:
                    print(f"WARNING: Failed to validate a random credential from manager by creating a temp client: {e_val}. App may rely on non-validated credentials.")
                    # Still return True if credentials exist, as the app might still function with other valid credentials.
                    # The per-request client creation will be the ultimate test for a specific credential.
                    return True # Credentials exist, even if one failed validation here.
            elif total_creds > 0 : # Credentials listed but get_random_credentials returned None
                 print(f"WARNING: {total_creds} credentials reported by manager, but could not retrieve one for validation. Problems might occur.")
                 return True # Still, credentials are listed.
            else: # No creds from get_random_credentials and total_creds is 0
                 print("ERROR: No credentials available after attempting to load from all sources.")
                 return False # No credentials reported by manager and get_random_credentials gave none.
        else:
            print("ERROR: Credential Manager reports no available credentials after processing all sources.")
            return False

    except Exception as e:
        print(f"CRITICAL ERROR during Vertex AI credential setup: {e}")
        return False