# This script initializes TeLLMgramBot with useful functions
import os
import sys

# Handle relative modules if running this script vs. elsewhere
try:
    from .utils import read_text, generate_file_path, read_yaml
except ImportError:
    from utils import read_text, generate_file_path, read_yaml

INIT_BOT_CONFIG = {
    'bot_owner': '<YOUR USERNAME>',
    'bot_nickname': 'Testy',
    'bot_initials': 'TB',
    'chat_model': 'gpt-5-mini',
    'url_model': 'gpt-5',
    'token_limit': None,
    'persona_temp': None,
    'persona_prompt': 'You are a generic test bot powered by a user-configured LLM.'
}

INIT_BOT_CONFIG_COMMENTS = {
    'token_limit': '# Optional, overrides the chat_model\'s default max token limit',
    'persona_temp': '# Optional, LLM temperature 0.0-2.0 (default: model\'s default)',
}

def execution_dir() -> str:
    """
    Initializes the top-level execution's path directory.
    """
    execution = os.path.abspath(sys.argv[0])
    if os.path.basename(execution) == 'initialize.py':
        execution = os.path.dirname(execution)
    return os.path.dirname(execution)

def init_directories():
    """
    Checks for TeLLMgramBot directories on configurations, prompts, conversations, and error logs.
    If each environment variable is not defined, set directory with base execution path:
    - TELLMGRAMBOT_CONFIGS_PATH
    - TELLMGRAMBOT_PROMPTS_PATH
    - TELLMGRAMBOT_ERRORLOGS_PATH
    - TELLMGRAMBOT_SESSIONLOGS_PATH
    """
    for directory in ['configs', 'prompts', 'errorlogs', 'sessionlogs']:
        env_var = f"TELLMGRAMBOT_{directory.upper()}_PATH"
        if not os.environ.get(env_var):
            os.environ[env_var] = os.path.join(execution_dir(), directory)
            try:
                if not os.path.exists(os.environ[env_var]):
                    os.makedirs(os.environ[env_var])
                    print(f"Created {env_var} directory: {os.environ[env_var]}")
            except Exception as e:
                sys.exit(f"{type(e).__name__} on {env_var} directory '{os.environ[env_var]}': {e}\nExiting...")

def _requires_provider(config_path: str) -> dict:
    """
    Read config.yaml and return which provider keys are required based on model prefixes.

    Provider inference mirrors factory.py: models starting with 'claude-' require Anthropic,
    all other non-empty models require OpenAI. If config.yaml is missing or unreadable,
    both providers are marked as required to err on the side of caution.

    Args:
        config_path: Path to config.yaml file.

    Returns:
        Dict with keys 'openai' and 'anthropic', each a bool indicating if that provider key is needed.
    """
    needs = {'openai': False, 'anthropic': False}
    try:
        config = read_yaml(config_path) or {}
        for field in ('chat_model', 'url_model'):
            model = config.get(field, '')
            if isinstance(model, str) and model.startswith('claude-'):
                needs['anthropic'] = True
            elif isinstance(model, str) and model:
                needs['openai'] = True
    except Exception:
        # If config is missing or unreadable, require both providers to be safe
        needs['openai'] = True
        needs['anthropic'] = True
    return needs

def init_keys():
    """
    TeLLMgramBot utilizes the following API keys via environment variables:
    - TELLMGRAMBOT_OPENAI_API_KEY     (required only when chat_model or url_model starts with 'gpt-' or other non-claude prefix)
    - TELLMGRAMBOT_ANTHROPIC_API_KEY  (required only when chat_model or url_model starts with 'claude-')
    - TELLMGRAMBOT_TELEGRAM_API_KEY   (always required)
    - TELLMGRAMBOT_VIRUSTOTAL_API_KEY (always required)

    Provider requirement is inferred from the configured models in config.yaml.
    If an environment variable is not set, its key file is read from the base execution
    path or TELLMGRAMBOT_KEYS_PATH. If no file exists, a placeholder is created.
    - openai.key
    - telegram.key
    - virustotal.key
    - anthropic.key
    """
    # Determine which provider keys are needed based on configured models
    config_path = os.path.join(os.environ.get('TELLMGRAMBOT_CONFIGS_PATH', execution_dir()), 'config.yaml')
    provider_needs = _requires_provider(config_path)

    # Build the ordered set of keys to check: provider keys first (conditional), then always-required keys
    provider_keys = {
        'openai': 'https://platform.openai.com/api-keys',
        'anthropic': 'https://docs.anthropic.com/en/api/getting-started',
    }
    always_keys = {
        'telegram': 'https://core.telegram.org/api',
        'virustotal': 'https://developers.virustotal.com/reference/overview'
    }

    keys = {k: v for k, v in provider_keys.items() if provider_needs.get(k)} | always_keys

    keys_need = 0
    for key, url in keys.items():
        env_var = f"TELLMGRAMBOT_{key.upper()}_API_KEY"
        if not os.environ.get(env_var):
            path = os.path.join(os.environ.get('TELLMGRAMBOT_KEYS_PATH', execution_dir()), f"{key}.key")
            if not os.path.exists(path):
                # Create a basic *.key file for user to input next run
                generate_file_path(os.path.dirname(path), os.path.basename(path), "API key",
                    f"YOUR {key.upper()} API KEY HERE - {url}\n"
                )
                keys_need = 1
                continue
            else:
                # Load each API key by file into its environment variable
                os.environ[env_var] = read_text(path)
                print(f"Loaded {env_var} secret from: {path}")

        # API keys must be defined and not have any spaces
        if not os.environ.get(env_var) or any(c.isspace() for c in os.environ.get(env_var)):
            print(f"{env_var} secret cannot be blank or have any whitelines")
            keys_need = 1

    # Stop program if one of the API keys is needed or has an error
    if keys_need:
        sys.exit("Need appropriate API keys to run TeLLMgramBot! Exiting...")

def init_bot_config(file='config.yaml') -> str:
    """
    Ensure bot configuration file (default 'config.yaml') is created.

    Creates a basic bot configuration file in TELLMGRAMBOT_CONFIGS_PATH with default values
    and inline parameter comments. Optional parameters (those with None values in INIT_BOT_CONFIG)
    are populated with descriptions from INIT_BOT_CONFIG_COMMENTS instead of values, helping
    users understand when to set them.

    Args:
        file: Name of the configuration file to create (default: 'config.yaml').

    Returns:
        Path to the created or existing configuration file.

    Raises:
        SystemExit: If TELLMGRAMBOT_CONFIGS_PATH environment variable is not defined,
                   or if file creation fails.
    """
    env_var = "TELLMGRAMBOT_CONFIGS_PATH"
    try:
        path = os.path.join(os.environ[env_var], file)
        if not os.path.exists(path):
            # Create a basic TeLLMgramBot configuration by filename
            generate_file_path(os.path.dirname(path), file, "bot configuration")
            with open(path, 'w') as f:
                # Add configuration with default values except prompt:
                for parameter, value in INIT_BOT_CONFIG.items():
                    if parameter == 'persona_prompt':
                        continue
                    else:  # Write parameter, with inline comment if no value
                        f.write('%s : %s\n' % (
                            parameter.ljust(12),
                            value if value else INIT_BOT_CONFIG_COMMENTS.get(parameter, '# Optional')
                        ))
        return path
    except KeyError:
        sys.exit(f"{env_var} must be defined to create file '{file}'! Exiting...")
    except Exception as e:
        sys.exit(f"{type(e).__name__} on {env_var} file '{file}': {e}\nExiting...")

def init_tokenGPT_config(file='tokenGPT.yaml') -> str:
    """
    Ensure TokenGPT configuration file (default 'tokenGPT.yaml') is created.
    Must specify the environment variable 'TELLMGRAMBOT_CONFIGS_PATH'.
    """
    env_var = "TELLMGRAMBOT_CONFIGS_PATH"
    try:
        # Create a basic GPT configuration of token parameters by filename
        return generate_file_path(os.environ[env_var], file, "TokenGPT configuration",
            "# Available parameters per model with default values:\n"
            "#   max_tokens: 4097\n"
            "#   tokens_per_message: 3  (OpenAI models only)\n"
            "#   tokens_per_name: 1     (OpenAI models only)\n"
            "# For OpenAI's updated list of models, please see:\n"
            "#   https://platform.openai.com/docs/models\n"
            "# For Anthropic's updated list of models, please see:\n"
            "#   https://docs.anthropic.com/en/docs/about-claude/models\n"
            "'gpt-4o':\n"
            "  max_tokens: 128000\n"
            "'gpt-4o-mini':\n"
            "  max_tokens: 128000\n"
            "'gpt-5':\n"
            "  max_tokens: 400000\n"
            "'gpt-5-mini':\n"
            "  max_tokens: 400000\n"
            "'gpt-5-nano':\n"
            "  max_tokens: 400000\n"
            "'claude-opus-4-6':\n"
            "  max_tokens: 200000\n"
            "'claude-sonnet-4-6':\n"
            "  max_tokens: 200000\n"
            "'claude-haiku-4-5':\n"
            "  max_tokens: 200000\n"
            "'claude-haiku-4-5-20251001':\n"
            "  max_tokens: 200000\n"
        )
    except KeyError:
        sys.exit(f"{env_var} must be defined to create file '{file}'! Exiting...")
    except Exception:
        sys.exit(f"Error while generating {env_var} file '{file}'! Exiting...")

def init_bot_prompt(file='test_personality.prmpt') -> str:
    """
    Ensure prompt file is created that defines a bot personality (or system content).
    Must specify the environment variable 'TELLMGRAMBOT_PROMPTS_PATH'.
    """
    env_var = "TELLMGRAMBOT_PROMPTS_PATH"
    try:
        return generate_file_path(os.environ[env_var], file, "bot personality prompt",
            "You are a generic test bot powered by a user-configured LLM. "
            "You respond helpfully to messages and, when a URL is provided in [square brackets], "
            "you fetch and analyze its contents after verifying it is safe via VirusTotal.\n"
        )
    except KeyError:
        sys.exit(f"{env_var} must be defined to create file '{file}'! Exiting...")
    except Exception:
        sys.exit(f"Error while generating {env_var} file '{file}'! Exiting...")

def init_url_prompt(file='url_analysis.prmpt') -> str:
    """
    Ensure prompt file is created that defines how the bot will analyze URLs via square brackets [].
    Must specify the environment variable 'TELLMGRAMBOT_PROMPTS_PATH'.
    """
    env_var = "TELLMGRAMBOT_PROMPTS_PATH"
    try:
        return generate_file_path(os.environ[env_var], file, "URL analysis prompt",
            "The user has provided a URL to perform some level of analysis. You will infer the nature of the "
            "analysis from the user's query.\n\n"
            "The contents of the URL mentioned have already been harvested and cleansed. Note the URL contents "
            "will likely have sections of text that are less relevant to the user's question (headers, footers, "
            "menus, ads, etc.). You will need to ignore those sections of text and focus on the main content of "
            "the page.\n\n"
            "The contents of the URL are shown below:\n"
            "BEGIN URL CONTENTS\n"
            "{url_content}\n"
            "END URL CONTENTS\n"
        )
    except KeyError:
        sys.exit(f"{env_var} must be defined to create file '{file}'! Exiting...")
    except Exception:
        sys.exit(f"Error while generating {env_var} file '{file}'! Exiting...")

def init_structure():
    """
    Performs the whole TeLLMgramBot first-run setup including:
    - Directories for configurations, prompts, and conversation/error logs.
    - Provider-conditional API keys (OpenAI/Anthropic determined by config.yaml model prefixes).
    - Configuration files (config.yaml, tokenGPT.yaml) with inline parameter descriptions.
    - System prompt files (test_personality.prmpt, url_analysis.prmpt).

    Provider key requirements are inferred from config.yaml model prefixes (claude-* → Anthropic,
    gpt-* → OpenAI, etc.) to streamline setup for single-provider deployments.
    """
    init_directories()
    init_keys()
    init_url_prompt()
    init_tokenGPT_config()

# Run this script if the first time setting up TeLLMgramBot
if __name__ == '__main__':
    init_structure()
    init_bot_config()
    init_bot_prompt()
    print("TeLLMgramBot setup complete!")
