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

# 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)',
}

# Append the framework-owned system appendix to the persona prompt.
# Teaches an LLM how cross-chat memory works without requiring persona authors to include it.
_SYSTEM_APPENDIX = (
    "You have persistent memory across conversations. Past interactions are selectively loaded into "
    "your current context window based on relevance and token budget. Treat anything visible in your "
    "current conversation history as reliable memory; do not deny any context visible to you.\n\n"
    "Cross-chat memory rules (apply only to what is present in your current context):\n"
    "- In a group chat: when the user's private chat history with you is included in your current context "
    "(cross-pollination), use it naturally. Do not deny having context that is visible to you.\n"
    "- In a private chat: if your current context includes shared group chat history that both you "
    "and the user are active in, use it to answer questions about those group conversations.\n"
    "- If a user has /private mode ON, their private-chat messages are excluded from group contexts. "
    "You may acknowledge this if relevant, but do not generalize it to deny all cross-chat context "
    "that is present in your current history.\n"
)

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: dict) -> dict:
    """
    Return which provider keys are required based on model prefixes in the parsed config dict.

    Provider inference mirrors factory.py: models starting with 'claude-' require Anthropic,
    all other non-empty models require OpenAI.

    Args:
        config: Parsed bot configuration dict (from init_bot_config or read_yaml on config.yaml).

    Returns:
        Dict with keys 'openai' and 'anthropic', each a bool indicating if that provider key is needed.
    """
    needs = {'openai': False, 'anthropic': False}
    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
    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(config: Optional[dict] = None) -> ApiKeyStatus:
    """
    Load API keys and return an 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. Key files are read from TELLMGRAMBOT_KEYS_PATH or the base execution path.
    Placeholder files are created for any missing keys. Prints API key status summary
    before returning.

    Args:
        config: Optional parsed bot configuration dict. If None, reads from
                TELLMGRAMBOT_CONFIGS_PATH/config.yaml (or execution_dir/config.yaml).

    Returns:
        ApiKeyStatus with chat_enabled and url_analysis_enabled flags set based on which keys are available.
    """
    # If no config provided, fall back to reading the default config file.
    # Failures here are non-fatal: config = {} means no models configured.
    if config is None:
        try:
            config_path = os.path.join(os.environ.get('TELLMGRAMBOT_CONFIGS_PATH', execution_dir()), 'config.yaml')
            config = read_yaml(config_path) or {}
        except Exception:
            config = {}

    # Determine which provider keys are needed and which feature each model uses.
    provider_needs = _requires_provider(config)
    chat_provider = _model_provider(config.get('chat_model', ''))
    url_provider = _model_provider(config.get('url_model', ''))

    # Build the full set of keys to check: only the required provider key(s) + 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

    available = set()  # tracks which keys loaded successfully
    for key, url in keys.items():
        env_var = f"TELLMGRAMBOT_{key.upper()}_API_KEY"

        # If the key isn't already in the environment, try loading it from a key file.
        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):
                # Key file missing: create a placeholder and warn (or exit for Telegram).
                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...")
                # Build a human-readable list of which features this key covers.
                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:
                # Key file exists: load it into the environment for this process.
                os.environ[env_var] = read_text(path)
                print(f"Loaded {env_var} secret from: {path}")

        # Reject blank or whitespace-only values (common copy-paste mistake).
        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)

    # Disable features whose required keys are not available
    key_status = ApiKeyStatus()

    # Chat requires the provider key matching chat_model (openai or anthropic).
    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"
        )

    # URL analysis requires both the provider key for url_model AND the VirusTotal 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)

    print_api_key_status(key_status)
    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: str = 'config.yaml') -> dict:
    """
    Returns contents of a bot configuration file (default 'config.yaml' if 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:
        Parsed configuration dict (empty dict if file exists but is blank).
    """
    env_var = "TELLMGRAMBOT_CONFIGS_PATH"
    try:
        text = ''.join(
            '%s : %s\n' % (
                parameter.ljust(12),
                value if value else INIT_BOT_CONFIG_COMMENTS.get(parameter, '# Optional')
            )
            for parameter, value in INIT_BOT_CONFIG.items()
            if parameter != 'persona_prompt'
        )
        return read_yaml(
            generate_file_path(os.environ[env_var], file, "bot configuration", text)
        ) or {}
    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: str = 'models.yaml') -> dict:
    """
    Ensure model token limit configuration file (default 'models.yaml') is created, and
    return its contents. Must specify the environment variable 'TELLMGRAMBOT_CONFIGS_PATH'.
    """
    env_var = "TELLMGRAMBOT_CONFIGS_PATH"
    try:
        return read_yaml(
            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"
            )
        ) or {}
    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_bot_prompt(file: str = 'test_personality.prmpt') -> str:
    """
    Ensure prompt file is created that defines a bot personality (or system content), and
    return its contents. Must specify the environment variable 'TELLMGRAMBOT_PROMPTS_PATH'.
    """
    env_var = "TELLMGRAMBOT_PROMPTS_PATH"
    try:
        return read_text(
            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 as e:
        sys.exit(f"{type(e).__name__} on {env_var} file '{file}': {e}\nExiting...")

def init_url_prompt(file: str = '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 read_text(
            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 as e:
        sys.exit(f"{type(e).__name__} on {env_var} file '{file}': {e}\nExiting...")

def init_structure(
    config_file: str = 'config.yaml',
    prompt_file: str = 'test_personality.prmpt'
) -> tuple[ApiKeyStatus, dict, str]:
    """
    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.

    Args:
        config_file: Name of the bot configuration file (default: 'config.yaml').
                     Resolved relative to TELLMGRAMBOT_CONFIGS_PATH.
        prompt_file: Name of the bot persona prompt file (default: 'test_personality.prmpt').
                     Resolved relative to TELLMGRAMBOT_PROMPTS_PATH.

    Returns:
        Tuple of (ApiKeyStatus, parsed config dict, persona prompt string with framework-owned
                  system appendix automatically appended).
    """
    init_directories()

    # Configurations for bot and LLM models
    config = init_bot_config(config_file)
    init_models_config()

    # System/persona prompts to use for LLM models
    prompt = init_bot_prompt(prompt_file)
    if not prompt.strip():
        prompt = INIT_BOT_CONFIG.get('persona_prompt', '')
        print(f"File '{prompt_file}' is empty, set default prompt '{prompt}'")
    prompt = prompt.rstrip() + "\n\n" + _SYSTEM_APPENDIX
    init_url_prompt()

    # Useful when setting a TeLLMgramBot object
    return (
        init_keys(config),
        config,
        prompt
    )

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