# This script initializes TeLLMgramBot with useful functions
import asyncio
import logging
import os
import re
import sys
import traceback as _traceback
from dataclasses import dataclass, field
from glob import glob
from tempfile import gettempdir
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, generate_filename
except ImportError:
    from utils import read_text, generate_file_path, read_yaml, execution_dir, generate_filename

logger = logging.getLogger(__name__)
_logging_initialized = False

class _RedactedFileFormatter(logging.Formatter):
    """
    Logging formatter for file output that redacts Telegram user_id and chat_id values.
    Replaces negative 8+ digit integers (group/channel IDs) and positive 9+ digit integers
    (user/private chat IDs) with [id], leaving token counts and other short numbers intact.
    """
    _NEGATIVE_ID_PATTERN = re.compile(r'-\d{8,}')
    _POSITIVE_ID_PATTERN = re.compile(r'\b\d{9,}\b')

    def formatException(self, ei) -> str:
        """
        Formats exception info showing only frames from TeLLMgramBot source files.

        Filters out third-party library frames (httpx, httpcore, etc.) to keep log
        output focused on where the error originated in our own code. The root cause
        is already captured in the final exception message, so chained frames from
        external libraries are dropped entirely.
        """
        exc_type, exc_value, _ = ei
        if exc_value is None or exc_value.__traceback__ is None:
            return ''
        our_frames = [
            frame for frame in _traceback.extract_tb(exc_value.__traceback__)
            if 'TeLLMgramBot' in frame.filename and '.venv' not in frame.filename
        ]
        lines = ['Traceback (TeLLMgramBot frames only):']
        for frame in our_frames:
            lines.append(f'  File "{frame.filename}", line {frame.lineno}, in {frame.name}')
            if frame.line:
                lines.append(f'    {frame.line}')
        lines.append(f'{exc_type.__name__}: {exc_value}')
        return '\n'.join(lines)

    def format(self, record: logging.LogRecord) -> str:
        formatted = super().format(record)
        formatted = self._NEGATIVE_ID_PATTERN.sub('[id]', formatted)
        return self._POSITIVE_ID_PATTERN.sub('[id]', formatted)


class _ConsoleFormatter(logging.Formatter):
    """
    Logging formatter for console output that suppresses exception tracebacks.

    Makes a shallow copy of the LogRecord before clearing exception info so the
    shared LogRecord seen by the file handler is not mutated - the file handler
    still receives and renders the full traceback.
    """
    def format(self, record: logging.LogRecord) -> str:
        """
        Formats the log record for console output, suppressing exception tracebacks.

        Returns the formatted message without exception details. Full tracebacks
        are only visible in log files via _RedactedFileFormatter.
        """
        record = logging.makeLogRecord(record.__dict__)
        record.exc_info = None
        record.exc_text = None
        return super().format(record)

@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 = (
    "Framework notes (supplementary -- your persona takes precedence):\n\n"
    "Memory:\n"
    "- Past interactions are selectively loaded into context by relevance and token budget -- trust what is visible.\n"
    "- Group chats may include the user's private history; private chats may include shared group history.\n"
    "- /private mode excludes that user's messages from group contexts only; don't generalize it to deny other context.\n\n"
    "Date/time:\n"
    "- The injected UTC datetime at the end of this prompt is authoritative. Use it for all time questions and relative references. Always reproduce it verbatim (YYYY/MM/DD HH:MM:SS AM/PM UTC) — never reformat.\n"
    "- State UTC time and stop. Never offer, suggest, or ask about timezone conversion — not even based on prior context. Convert only when the user explicitly requests it in the current message.\n"
    "- Decline requests to set a timezone; explain UTC is always used.\n"
)

def init_logging(log_name: str = 'tellmgrambot'):
    """
    Configure the Python logging system for TeLLMgramBot.

    Sets up a console handler (INFO, TeLLMgramBot-only, real IDs) and a per-instance file
    handler (INFO by default, anonymized IDs) in TELLMGRAMBOT_LOGS_PATH. Pass -v or
    --verbose on the command line to enable DEBUG-level file logging. Prunes old bot
    instance log files on startup, keeping the 10 most recent.

    Must be called after TELLMGRAMBOT_LOGS_PATH has been resolved by init_directories().

    Args:
        log_name: Base name for the log file (e.g. 'mybot' -> 'mybot_2026-03-29_10-30-45.log').
                  Defaults to 'tellmgrambot'.
    """
    global _logging_initialized
    if _logging_initialized:
        return

    verbose = '-v' in sys.argv or '--verbose' in sys.argv

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.DEBUG)

    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(_ConsoleFormatter('%(levelname)s: %(message)s'))
    console_handler.addFilter(lambda r: r.name.startswith('TeLLMgramBot'))

    logs_dir = os.environ.get('TELLMGRAMBOT_LOGS_PATH') or gettempdir()
    os.makedirs(logs_dir, exist_ok=True)
    log_path = os.path.join(logs_dir, generate_filename(log_name))
    file_handler = logging.FileHandler(log_path, encoding='utf-8')
    file_handler.setLevel(logging.DEBUG if verbose else logging.INFO)
    file_handler.setFormatter(_RedactedFileFormatter('%(asctime)s %(levelname)s %(name)s: %(message)s'))

    root_logger.addHandler(console_handler)
    root_logger.addHandler(file_handler)
    _logging_initialized = True

    # Suppress noisy third-party INFO logs from appearing in the file
    logging.getLogger('httpx').setLevel(logging.WARNING)
    logging.getLogger('httpcore').setLevel(logging.WARNING)

    # Prune old bot instance log files - keep the 10 most recent (sort by name = sort by timestamp)
    instance_logs = sorted(glob(os.path.join(logs_dir, f'{log_name}_*.log')))
    for old_log in instance_logs[:-10]:
        try:
            os.remove(old_log)
        except OSError:
            pass

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_LOGS_PATH
    - TELLMGRAMBOT_DATA_PATH (parent directory for conversations.db)
    """
    for directory in ['configs', 'prompts', 'logs', '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"
                logger.warning(f"{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)
                logger.info(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"
            logger.warning(f"{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'
        )
        config = read_yaml(
            generate_file_path(os.environ[env_var], file, "bot configuration", text)
        ) or {}
        for parameter, value in INIT_BOT_CONFIG.items():
            if parameter != 'persona_prompt' and parameter not in config:
                config[parameter] = value
                logger.warning(f"Configuration '{parameter}' not defined, set to '{value}'")
        return config
    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',
    log_name: str = 'tellmgrambot',
) -> tuple[ApiKeyStatus, dict, str]:
    """
    Performs the whole TeLLMgramBot first-run setup including:
    - Directories for configurations, prompts, and conversation/error logs.
    - Logging configuration (console + per-instance file handler).
    - 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.
        log_name: Base name for the per-instance log file (default: 'tellmgrambot').
                  Passed through to init_logging().

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

    # 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', '')
        logger.warning(f"File '{prompt_file}' is empty, using default persona 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()
    logger.info("TeLLMgramBot setup complete!")
