# This script initializes TeLLMgramBot with useful functions
import asyncio
import os
import sys
from dataclasses import dataclass, field

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

@dataclass
class ApiKeyStatus:
    """
    Tracks API key presence at startup and which features are enabled as a result.
    Used by TelegramBot to gate features gracefully rather than hard-failing.

    Attributes:
        chat_enabled: True if chat_model is configured and its provider key is available.
        url_analysis_enabled: True if url_model is configured and both its provider key and VirusTotal key are available.
        disabled_reasons: Dict mapping feature names to reason strings (e.g., {'chat': 'missing openai LLM API key'}).
    """
    chat_enabled: bool = True
    url_analysis_enabled: bool = True
    disabled_reasons: dict = field(default_factory=dict)

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 init_directories():
    """
    Checks for TeLLMgramBot directories and creates them if missing.
    If each environment variable is not defined, set directory with base execution path:
    - TELLMGRAMBOT_CONFIGS_PATH
    - TELLMGRAMBOT_PROMPTS_PATH
    - TELLMGRAMBOT_ERRORLOGS_PATH
    - TELLMGRAMBOT_DATA_PATH (parent directory for conversations.db)
    """
    for directory in ['configs', 'prompts', 'errorlogs', 'data']:
        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...")

    # Initialize the SQLite schema; mirror TelegramBot event loop pattern
    try:
        from .database import init_db
    except ImportError:
        from database import init_db
    try:
        loop = asyncio.get_running_loop()
    except RuntimeError:
        asyncio.run(init_db())
    else:
        loop.create_task(init_db())

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 _model_provider(model: str) -> str | None:
    """Return the provider name for a given model string, or None if model is empty."""
    if not isinstance(model, str) or not model:
        return None
    return 'anthropic' if model.startswith('claude-') else 'openai'

def init_keys() -> ApiKeyStatus:
    """
    Load API keys and return a ApiKeyStatus indicating which features are available.

    Telegram key is always required — bot exits if missing or invalid.
    All other keys (OpenAI, Anthropic, VirusTotal) are optional: missing keys
    disable the associated features rather than exiting.

    Provider keys (OpenAI/Anthropic) are loaded conditionally based on configured
    models in config.yaml. Key files are read from TELLMGRAMBOT_KEYS_PATH or the
    base execution path. Placeholder files are created for any missing keys.

    Returns:
        ApiKeyStatus with chat_enabled and url_analysis_enabled flags set based on which keys are available.
    """
    config_path = os.path.join(os.environ.get('TELLMGRAMBOT_CONFIGS_PATH', execution_dir()), 'config.yaml')
    provider_needs = _requires_provider(config_path)

    # Read per-feature provider requirements from config
    try:
        config = read_yaml(config_path) or {}
    except Exception:
        config = {}
    chat_provider = _model_provider(config.get('chat_model', ''))
    url_provider = _model_provider(config.get('url_model', ''))

    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

    available = set()  # tracks which keys loaded successfully
    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):
                generate_file_path(os.path.dirname(path), os.path.basename(path), "API key",
                    f"YOUR {key.upper()} API KEY HERE - {url}\n"
                )
                if key == 'telegram':
                    sys.exit("Telegram API key is required to run TeLLMgramBot! Exiting...")
                which = [f"chat_model ({config.get('chat_model', '?')})" if chat_provider == key else None,
                         f"url_model ({config.get('url_model', '?')})" if url_provider == key else None,
                         "VirusTotal integration" if key == 'virustotal' else None]
                reason = ", ".join(w for w in which if w) or "related features"
                print(f"Warning: {key}.key not found — {reason} will be disabled.")
                continue
            else:
                os.environ[env_var] = read_text(path)
                print(f"Loaded {env_var} secret from: {path}")

        # Keys must be defined and not contain whitespace
        if not os.environ.get(env_var) or any(c.isspace() for c in os.environ.get(env_var)):
            if key == 'telegram':
                sys.exit("Telegram API key is blank or invalid! Exiting...")
            which = [f"chat_model ({config.get('chat_model', '?')})" if chat_provider == key else None,
                     f"url_model ({config.get('url_model', '?')})" if url_provider == key else None,
                     "VirusTotal integration" if key == 'virustotal' else None]
            reason = ", ".join(w for w in which if w) or "related features"
            print(f"Warning: {env_var} is blank or invalid — {reason} will be disabled.")
            continue

        available.add(key)

    # Build API key status based on which keys are available
    key_status = ApiKeyStatus()

    chat_provider_ok = bool(chat_provider) and chat_provider in available
    if not chat_provider_ok:
        key_status.chat_enabled = False
        key_status.disabled_reasons['chat'] = (
            "chat_model not configured" if not chat_provider
            else f"missing {chat_provider} LLM API key"
        )

    virustotal_ok = 'virustotal' in available
    url_provider_ok = bool(url_provider) and url_provider in available
    if not url_provider_ok or not virustotal_ok:
        key_status.url_analysis_enabled = False
        reasons = []
        if not url_provider_ok:
            reasons.append("url_model not configured" if not url_provider else f"missing {url_provider} LLM API key")
        if not virustotal_ok:
            reasons.append("missing VirusTotal API key")
        key_status.disabled_reasons['url_analysis'] = ", ".join(reasons)

    return key_status

def print_api_key_status(key_status: ApiKeyStatus):
    """Print a startup summary of which API keys are available and which features they enable."""
    print("--- TeLLMgramBot API Key Status ---")
    features = [
        ('chat', 'Chat (LLM)', key_status.chat_enabled),
        ('url_analysis', 'URL Analysis', key_status.url_analysis_enabled),
    ]
    for key, label, enabled in features:
        if enabled:
            print(f"  {label:<12} : ENABLED")
        else:
            reason = key_status.disabled_reasons.get(key, 'unavailable')
            print(f"  {label:<12} : DISABLED ({reason})")
    print("-----------------------------------")

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_models_config(file='models.yaml') -> str:
    """
    Ensure model token limit configuration file (default 'models.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, "model token limit 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() -> ApiKeyStatus:
    """
    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, models.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.

    Returns:
        ApiKeyStatus indicating which features are available based on loaded API keys.
    """
    init_directories()
    key_status = init_keys()
    print_api_key_status(key_status)
    init_url_prompt()
    init_models_config()
    return key_status

# 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!")
