from __future__ import annotations
import os

import v0.relationalai as rai
from v0.relationalai.clients import config as cfg

def create_engine(engine_name: str, size: str,  use_direct_access=False):
    print('create_engine: about to call make_config')
    config = make_config(engine_name, use_direct_access=use_direct_access)

    provider = rai.Resources(config=config)
    print(f"Creating engine {engine_name}")
    provider.create_engine(name=engine_name, type="LOGIC", size=size)
    print(f"Engine {engine_name} created")

def delete_engine(engine_name: str, use_direct_access=False):
    print(f"Deleting engine {engine_name}")
    config = make_config(engine_name, use_direct_access=use_direct_access)
    provider = rai.Resources(config=config)
    provider.delete_engine(engine_name, "LOGIC")
    print(f"Engine {engine_name} deleted")

def make_config(engine_name: str | None = None, fetch_profile: bool = True, use_direct_access = False, show_full_traces = True, show_debug_logs = True, reuse_model = False) -> cfg.Config:
    # First try to load from raiconfig.toml
    try:
        config = cfg.Config()
        if config.file_path is not None:
            # Set test defaults
            config.set("reuse_model", reuse_model)
            rai.Resources(config=config)
            return config
    except Exception as e:
        print(f"Could not load config from file: {e}, trying environment variables")
    # If that fails, construct from environment variables
    cloud_provider = os.getenv("RAI_CLOUD_PROVIDER")

    if cloud_provider is None:
        cloud_provider = os.getenv("RAI_CLOUD_PROVIDER")

    print('cloud provider:', cloud_provider)

    if cloud_provider is None:
        raise ValueError("RAI_CLOUD_PROVIDER must be set")
    elif cloud_provider == "snowflake":
        sf_username = os.getenv("SF_TEST_ACCOUNT_USERNAME")
        sf_password = os.getenv("SF_TEST_ACCOUNT_PASSWORD")
        sf_account = os.getenv("SF_TEST_ACCOUNT_NAME")
        sf_role = os.getenv("SF_TEST_ROLE_NAME", "RAI_USER")
        sf_warehouse = os.getenv("SF_TEST_WAREHOUSE_NAME")
        sf_app_name = os.getenv("SF_TEST_APP_NAME")
        if sf_username is None or sf_password is None:
            raise ValueError(
                "SF_TEST_ACCOUNT_USERNAME, SF_TEST_ACCOUNT_PASSWORD, SF_TEST_ACCOUNT_NAME must be set if RAI_CLOUD_PROVIDER is set to 'snowflake'"
            )

        current_config = {
            "platform": "snowflake",
            "user": sf_username,
            "password": sf_password,
            "account": sf_account,
            "role": sf_role,
            "warehouse": sf_warehouse,
            "rai_app_name": sf_app_name,
            "use_direct_access": use_direct_access,
            "show_full_traces": show_full_traces,
            "show_debug_logs": show_debug_logs,
            "reuse_model": reuse_model,
        }
        if engine_name:
            current_config["engine"] = engine_name

        # For direct access, we use key-pair as the default for testing and need to configure additional parameters.
        if use_direct_access:
            authenticator=os.getenv("AUTHENTICATOR","")
            private_key_file=os.getenv("PRIVATE_KEY_FILE","")
            private_key_file_pwd=os.getenv("SF_TEST_ACCOUNT_KEY_PASSPHRASE","")

            current_config["authenticator"] = authenticator
            current_config["private_key_file"] = private_key_file
            current_config["private_key_file_pwd"]=private_key_file_pwd

        return cfg.Config(current_config, fetch=fetch_profile)

    else:
        raise ValueError(f"Unsupported cloud provider: {cloud_provider}")
