#!/usr/bin/env python3
# -*- coding: utf-8 -*-
#
# openmamba bot for Telegram
#
# Copyright (C) 2016-2026 by Silvan Calarco <silvan.calarco@mambasoft.it>
#
# GPL v3 license

from telegram import Update
from telegram.constants import ParseMode
from telegram.ext import (Application, CommandHandler, MessageHandler,
                          filters, ContextTypes)

import asyncio
import json
import logging
import os
import random
import time
import urllib.request
import urllib.parse

try:
    import anthropic
    ANTHROPIC_AVAILABLE = True
except ImportError:
    ANTHROPIC_AVAILABLE = False

# Enable logging
logging.basicConfig(format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
                    level=logging.INFO)

logger = logging.getLogger(__name__)
logging.getLogger('httpx').setLevel(logging.WARNING)
social_log_last_ids = dict()
conversation_history = dict()
pending_queries = {}    # {chat_id: {msg_id: asyncio.Task}}
admin_last_seen = {}    # {chat_id: float} timestamp of last admin/openmamba message
_admin_cache = {}       # {chat_id: (timestamp, frozenset of user_ids)}

MAX_HISTORY = 20
QUERY_DELAY_MIN = 30    # seconds before responding in group chats
QUERY_DELAY_MAX = 60
ADMIN_SUPPRESS_SECS = 600   # seconds to stay silent after admin posts
ADMIN_CACHE_TTL = 300       # seconds to cache admin list
GROUP_ANONYMOUS_BOT_ID = 1087968824  # Telegram's fixed ID for anonymous admin messages

AI_SYSTEM_PROMPTS = {
    'en': (
        "You are a helpful assistant for the openmamba GNU/Linux distribution. "
        "You help users with questions about packages, installation, configuration, "
        "build system (autodist, autospec, distromatic), and the openmamba ecosystem in general. "
        "You have access to tools that query the live openmamba package API: use them whenever "
        "the user asks about specific packages, repositories, dependencies, or build problems. "
        "Be concise and friendly. "
        "Do not use any markdown formatting in your responses: no **bold**, no _italic_, "
        "no ## headings, no bullet lists with *, no backticks. Use plain text only."
    ),
    'it': (
        "Sei un assistente per la distribuzione GNU/Linux openmamba. "
        "Aiuti gli utenti con domande su pacchetti, installazione, configurazione, "
        "sistema di build (autodist, autospec, distromatic) e l'ecosistema openmamba in generale. "
        "Hai accesso a strumenti che interrogano l'API live dei pacchetti openmamba: usali ogni volta "
        "che l'utente chiede di pacchetti specifici, repository, dipendenze o problemi di build. "
        "Sii conciso e cordiale. "
        "Non usare nessuna formattazione markdown nelle risposte: niente **grassetto**, niente _corsivo_, "
        "niente ## titoli, niente liste con *, niente backtick. Usa solo testo normale."
    ),
}

CONFIG_FILE = '/etc/openmamba-telegram-bot/main.conf'
SKILLS_DIR = '/etc/openmamba-telegram-bot/skills'


def read_config(path):
    config = {}
    try:
        with open(path) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#') or '=' not in line:
                    continue
                key, _, value = line.partition('=')
                value = value.strip()
                if len(value) >= 2 and value[0] == value[-1] and value[0] in ('"', "'"):
                    value = value[1:-1]
                config[key.strip()] = value
    except OSError:
        pass
    return config


def get_config():
    return read_config(CONFIG_FILE)


def get_api_url():
    return get_config().get('DISTROQUERY_API_URL', '')


def get_api_token():
    return get_config().get('DISTROQUERY_API_TOKEN', '')


def get_repositories():
    reps_str = get_config().get('AUTODIST_REPOSITORIES', '')
    return reps_str.strip('()').split()


def get_anthropic_key():
    return get_config().get('ANTHROPIC_API_KEY', '')


def get_chat_language(chat_id):
    """Return 'it' or 'en' for a given chat_id, based on TELEGRAM_AI_CHATS config.

    Config format (space-separated): -1001234567890:it -1009876543210:en
    Defaults to 'en' if the chat is not listed.
    """
    chats_str = get_config().get('TELEGRAM_AI_CHATS', '')
    for entry in chats_str.split():
        entry = entry.strip()
        if ':' not in entry:
            continue
        cid, lang = entry.split(':', 1)
        try:
            if int(cid) == chat_id:
                return lang
        except ValueError:
            continue
    return 'en'


async def get_admin_ids(context, chat_id):
    """Return the set of admin user IDs for a chat, with caching."""
    cached = _admin_cache.get(chat_id)
    if cached and time.time() - cached[0] < ADMIN_CACHE_TTL:
        return cached[1]
    try:
        admins = await context.bot.get_chat_administrators(chat_id)
        ids = frozenset(a.user.id for a in admins)
        _admin_cache[chat_id] = (time.time(), ids)
        return ids
    except Exception as e:
        logger.warning("Failed to get admin list for chat %s: %s" % (chat_id, e))
        return frozenset()


def load_skills(skills_dir=SKILLS_DIR):
    """Load all skill files from skills_dir. Each file may have a frontmatter header:
        ---
        name: skill name
        keywords: word1, word2, word3
        ---
        skill content...
    Returns a list of dicts with keys: name, keywords (list), content (str).
    """
    skills = []
    try:
        entries = os.listdir(skills_dir)
    except OSError:
        return skills
    for fname in sorted(entries):
        if not fname.endswith('.md'):
            continue
        fpath = os.path.join(skills_dir, fname)
        try:
            with open(fpath, encoding='utf-8') as f:
                raw = f.read()
        except OSError:
            continue
        skill = {'name': fname, 'keywords': [], 'content': raw}
        if raw.startswith('---'):
            parts = raw.split('---', 2)
            if len(parts) >= 3:
                header, body = parts[1], parts[2].strip()
                skill['content'] = body
                for line in header.splitlines():
                    line = line.strip()
                    if line.startswith('name:'):
                        skill['name'] = line[5:].strip()
                    elif line.startswith('keywords:'):
                        kw_str = line[9:].strip()
                        skill['keywords'] = [k.strip().lower() for k in kw_str.split(',') if k.strip()]
        skills.append(skill)
    return skills


def select_skills(question, skills):
    """Return skills whose keywords appear in the question (case-insensitive)."""
    q = question.lower()
    return [s for s in skills if any(kw in q for kw in s['keywords'])]


def fetch_json(url, token=None):
    headers = {'Accept': 'application/json'}
    if token:
        headers['Authorization'] = f'Bearer {token}'
    req = urllib.request.Request(url, headers=headers)
    try:
        with urllib.request.urlopen(req, timeout=10) as resp:
            return json.load(resp)
    except Exception as e:
        logger.error("API request failed for %s: %s" % (url, e))
        return None


TOOLS = [
    {
        "name": "search_packages",
        "description": "Search packages across all openmamba repositories by name, summary or description.",
        "input_schema": {
            "type": "object",
            "properties": {
                "q": {"type": "string", "description": "Search term"},
                "arch": {"type": "string", "description": "Comma-separated architectures (e.g. src,x86_64,aarch64). Default: all"},
                "per_page": {"type": "integer", "description": "Results per page (default 20)"},
            },
            "required": ["q"],
        },
    },
    {
        "name": "list_repositories",
        "description": "List all openmamba repositories with their architectures and descriptions.",
        "input_schema": {"type": "object", "properties": {}},
    },
    {
        "name": "get_package_details",
        "description": (
            "Get details of a source package in a specific repository, "
            "including version, description, binary packages, and changelog."
        ),
        "input_schema": {
            "type": "object",
            "properties": {
                "repository": {"type": "string", "description": "Repository tag (e.g. base, devel-autodist)"},
                "package": {"type": "string", "description": "Source package name"},
            },
            "required": ["repository", "package"],
        },
    },
    {
        "name": "get_binary_package_details",
        "description": "Get details of a binary package (provides, requires, files) for a specific architecture.",
        "input_schema": {
            "type": "object",
            "properties": {
                "repository": {"type": "string"},
                "package": {"type": "string"},
                "arch": {"type": "string", "description": "Architecture (e.g. x86_64, aarch64, i586)"},
            },
            "required": ["repository", "package", "arch"],
        },
    },
    {
        "name": "get_repository_problems",
        "description": "List warnings and packages needing rebuild in a repository.",
        "input_schema": {
            "type": "object",
            "properties": {
                "repository": {"type": "string", "description": "Repository tag"},
            },
            "required": ["repository"],
        },
    },
    {
        "name": "find_providers",
        "description": "Find packages that provide a given capability or dependency in a repository.",
        "input_schema": {
            "type": "object",
            "properties": {
                "repository": {"type": "string"},
                "requirement": {"type": "string", "description": "Capability name to search for"},
            },
            "required": ["repository", "requirement"],
        },
    },
    {
        "name": "find_file_providers",
        "description": "Find packages that provide a specific file path in a repository.",
        "input_schema": {
            "type": "object",
            "properties": {
                "repository": {"type": "string"},
                "filepath": {"type": "string", "description": "File path (e.g. /bin/bash)"},
                "arch": {"type": "string", "description": "Architecture (e.g. x86_64, aarch64)"},
            },
            "required": ["repository", "filepath", "arch"],
        },
    },
]


def execute_tool(name, inputs, api_url):
    """Dispatch a tool call and return the result as a JSON string."""
    q = inputs.get
    result = None
    if name == 'search_packages':
        params = {'q': q('q'), 'per_page': q('per_page', 20)}
        if q('arch'):
            params['arch'] = q('arch')
        result = fetch_json(f"{api_url}/search?{urllib.parse.urlencode(params)}")
    elif name == 'list_repositories':
        result = fetch_json(f"{api_url}/repositories")
    elif name == 'get_package_details':
        result = fetch_json(
            f"{api_url}/package/{urllib.parse.quote(q('repository'))}/{urllib.parse.quote(q('package'))}"
        )
    elif name == 'get_binary_package_details':
        result = fetch_json(
            f"{api_url}/package/{urllib.parse.quote(q('repository'))}"
            f"/{urllib.parse.quote(q('package'))}/{urllib.parse.quote(q('arch'))}"
        )
    elif name == 'get_repository_problems':
        result = fetch_json(f"{api_url}/problems/{urllib.parse.quote(q('repository'))}")
    elif name == 'find_providers':
        result = fetch_json(
            f"{api_url}/providers/{urllib.parse.quote(q('repository'))}/{urllib.parse.quote(q('requirement'))}"
        )
    elif name == 'find_file_providers':
        result = fetch_json(
            f"{api_url}/file_providers/{urllib.parse.quote(q('repository'))}"
            f"/{urllib.parse.quote(q('filepath'), safe='')}/{urllib.parse.quote(q('arch'))}"
        )
    if result is None:
        return '{"error": "tool call failed or API not available"}'
    return json.dumps(result)


HELP_TEXT = {
    'en': (
        'Send any text to ask the AI assistant about openmamba.\n'
        'Send /search [term] to search for packages.\n'
        'Send /details [name] to see details of a source package.\n'
        'Send /ask [question] to ask the AI assistant explicitly.\n'
        'Send /reset to clear the AI conversation history.\n'
        'Send /set [seconds] to enable notifications.\n'
        'Send /unset to disable notifications.'
    ),
    'it': (
        'Invia del testo per chiedere all\'assistente AI su openmamba.\n'
        'Invia /search [termine] per cercare pacchetti.\n'
        'Invia /details [nome] per vedere i dettagli di un pacchetto sorgente.\n'
        'Invia /ask [domanda] per chiedere all\'assistente AI esplicitamente.\n'
        'Invia /reset per cancellare la cronologia della conversazione AI.\n'
        'Invia /set [secondi] per abilitare le notifiche.\n'
        'Invia /unset per disabilitare le notifiche.'
    ),
}


async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
    lang = get_chat_language(update.effective_chat.id)
    prefix = {
        'en': 'Hi! This is the openmamba Linux Bot.\n',
        'it': 'Ciao! Questo è il Bot di openmamba Linux.\n',
    }[lang]
    await update.message.reply_text(prefix + HELP_TEXT[lang])


async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
    lang = get_chat_language(update.effective_chat.id)
    await update.message.reply_text(HELP_TEXT[lang])


async def search_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handle /search <term>: search packages by name/keyword."""
    user = update.message.from_user
    if not context.args:
        await update.message.reply_text('Usage: /search <term>')
        return
    query_text = ' '.join(context.args)
    logger.info("Search by %s (chat %s): %s" % (
        user.first_name, update.effective_chat.id, query_text))

    api_url = get_api_url()
    if not api_url:
        await update.message.reply_text('API not configured.')
        return

    params = urllib.parse.urlencode({'q': query_text, 'per_page': 20, 'arch': 'src'})
    d = fetch_json(f"{api_url}/search?{params}")
    if not d:
        await update.message.reply_text('Search failed.')
        return

    packages = d.get('packages', [])
    if not packages:
        await update.message.reply_text('No results found.')
        return

    response = ""
    for pkg in packages:
        response += "<b>%s</b> %s-%s (%s)\n%s\n\n" % (
            pkg.get('name', ''), pkg.get('version', ''), pkg.get('release', ''),
            pkg.get('repository', ''), pkg.get('summary', ''))

    await update.message.reply_text(response[:4096], parse_mode=ParseMode.HTML)


async def _do_ask(update: Update, context: ContextTypes.DEFAULT_TYPE, question: str):
    """Call the AI assistant with question and send the reply."""
    lang = get_chat_language(update.effective_chat.id)
    unavailable = {'en': 'AI assistant temporarily unavailable.',
                   'it': 'Assistente AI temporaneamente non disponibile.'}
    not_configured = {'en': 'AI assistant not configured.',
                      'it': 'Assistente AI non configurato.'}

    if not ANTHROPIC_AVAILABLE:
        await update.message.reply_text(not_configured[lang])
        return

    api_key = get_anthropic_key()
    if not api_key:
        await update.message.reply_text(not_configured[lang])
        return

    chat_id = update.effective_chat.id
    user = update.message.from_user
    logger.info("AI ask from %s (chat %s): %s" % (user.first_name, chat_id, question))

    history = conversation_history.setdefault(chat_id, [])

    try:
        skills = load_skills()
        matched = select_skills(question, skills)
        system_prompt = AI_SYSTEM_PROMPTS[lang]
        if matched:
            skills_text = '\n\n'.join(s['content'] for s in matched)
            system_prompt += '\n\n' + skills_text
            logger.info("Skills matched for chat %s: %s" % (chat_id, [s['name'] for s in matched]))

        api_url = get_api_url()
        client = anthropic.Anthropic(api_key=api_key)
        messages = list(history) + [{"role": "user", "content": question}]

        while True:
            response = client.messages.create(
                model="claude-sonnet-4-6",
                max_tokens=1024,
                system=system_prompt,
                messages=messages,
                tools=TOOLS if api_url else [],
            )
            if response.stop_reason == 'tool_use':
                messages.append({"role": "assistant", "content": response.content})
                tool_results = []
                for block in response.content:
                    if block.type == 'tool_use':
                        logger.info("Tool call: %s(%s)" % (block.name, block.input))
                        result = execute_tool(block.name, block.input, api_url)
                        tool_results.append({
                            "type": "tool_result",
                            "tool_use_id": block.id,
                            "content": result,
                        })
                messages.append({"role": "user", "content": tool_results})
            else:
                answer = response.content[0].text
                break

        history.append({"role": "user", "content": question})
        history.append({"role": "assistant", "content": answer})
        if len(history) > MAX_HISTORY:
            conversation_history[chat_id] = history[-MAX_HISTORY:]

        await update.message.reply_text(answer[:4096])
    except Exception as e:
        logger.error("Claude API error: %s" % e)
        await update.message.reply_text(unavailable[lang])


async def _delayed_ask(update: Update, context: ContextTypes.DEFAULT_TYPE, delay: float):
    """Wait delay seconds, then run the AI assistant if not cancelled."""
    try:
        await asyncio.sleep(delay)
    except asyncio.CancelledError:
        return
    chat_id = update.effective_chat.id
    msg_id = update.message.message_id
    pending_queries.get(chat_id, {}).pop(msg_id, None)
    await _do_ask(update, context, update.message.text)


def _is_bot_mentioned(update: Update, context: ContextTypes.DEFAULT_TYPE) -> bool:
    """Return True if the bot is @mentioned in the message."""
    bot_username = context.bot.username
    if not bot_username:
        return False
    for entity in (update.message.entities or []):
        if entity.type == 'mention':
            name = update.message.text[entity.offset:entity.offset + entity.length]
            if name.lstrip('@').lower() == bot_username.lower():
                return True
    return False


async def query(update: Update, context: ContextTypes.DEFAULT_TYPE):
    chat_id = update.effective_chat.id
    msg_id = update.message.message_id
    sender = update.message.from_user
    chat_type = update.effective_chat.type

    if chat_type not in ('group', 'supergroup'):
        await _do_query(update, context)
        return

    # In group chats: check if sender is openmamba or a chat admin
    admin_ids = await get_admin_ids(context, chat_id)
    if (sender.username == 'openmamba') or (sender.id in admin_ids) or (sender.id == GROUP_ANONYMOUS_BOT_ID):
        # Cancel any pending bot responses and record admin activity
        for task in list(pending_queries.get(chat_id, {}).values()):
            task.cancel()
        pending_queries[chat_id] = {}
        admin_last_seen[chat_id] = time.time()
        logger.info("Admin activity in chat %s by %s, pending queries cancelled" % (
            chat_id, sender.username or sender.id))
        return

    # Cancel pending response if this message is an explicit reply to one
    if update.message.reply_to_message:
        replied_id = update.message.reply_to_message.message_id
        task = pending_queries.get(chat_id, {}).pop(replied_id, None)
        if task:
            task.cancel()
            logger.info("Pending query %s cancelled: reply received" % replied_id)

    # Respond immediately if bot is explicitly @mentioned
    if _is_bot_mentioned(update, context):
        await _do_ask(update, context, update.message.text)
        return

    # Suppress if admin was recently active
    if time.time() - admin_last_seen.get(chat_id, 0) < ADMIN_SUPPRESS_SECS:
        logger.info("Query suppressed in chat %s: admin recently active" % chat_id)
        return

    # Schedule delayed AI response
    delay = random.uniform(QUERY_DELAY_MIN, QUERY_DELAY_MAX)
    task = asyncio.create_task(_delayed_ask(update, context, delay))
    pending_queries.setdefault(chat_id, {})[msg_id] = task
    logger.info("Query from %s scheduled in %.0fs (chat %s)" % (
        sender.first_name, delay, chat_id))


async def details(update: Update, context: ContextTypes.DEFAULT_TYPE):
    user = update.message.from_user
    logger.info("Details of %s: %s" % (user.first_name, context.args))

    api_url = get_api_url()
    if not api_url or not context.args:
        await update.message.reply_text('Usage: /details <package>')
        return

    pkg_name = context.args[0]
    response = ""
    for rep in get_repositories():
        d = fetch_json(f"{api_url}/package/{urllib.parse.quote(rep)}/{urllib.parse.quote(pkg_name)}")
        if not d or 'error' in d:
            continue
        version = d.get('version', '')
        release = d.get('release', '')
        summary = d.get('summary', '')
        url = d.get('url', '')
        description = d.get('description', '')
        response += "<b>%s</b> %s-%s (%s)\n%s\n%s\n\n<i>%s</i>\n\n" % (
            pkg_name, version, release, rep, summary, url, description)
        archs_data = d.get('children', {}).get('archs', {})
        for arch, pkgs in archs_data.items():
            for p in pkgs:
                response += "<b>%s</b>(%s) " % (p.get('name', ''), arch)
        if archs_data:
            response += "\n\n\n"

    if response:
        await update.message.reply_text(response[:4096], parse_mode=ParseMode.HTML)
    else:
        await update.message.reply_text('No results found.')


async def alarm(context: ContextTypes.DEFAULT_TYPE):
    chat_id = context.job.chat_id
    api_url = get_api_url()
    token = get_api_token()
    if not api_url or not token:
        return

    last_id = social_log_last_ids.get(chat_id, 0)
    params = urllib.parse.urlencode({'from_id': last_id, 'limit': 100})
    d = fetch_json(f"{api_url}/social_log?{params}", token)
    if not d:
        return

    entries = d if isinstance(d, list) else d.get('entries', [])
    response = ""
    last_seen = last_id
    for e in entries:
        eid = e.get('id', 0)
        if eid <= last_id:
            continue
        user = e.get('user', '')
        text = e.get('text', '')
        time = e.get('time', '')
        etype = e.get('type', '')
        if etype == 'job':
            response += "Job run by <i>%s</i> %s (%s)\n" % (user, text, time)
        else:
            response += "<i>%s</i> %s (%s)\n" % (user, text, time)
        if eid > last_seen:
            last_seen = eid
    if last_seen > last_id:
        social_log_last_ids[chat_id] = last_seen
    if response:
        await context.bot.send_message(chat_id, response, parse_mode=ParseMode.HTML)


async def set_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
    chat_id = update.message.chat_id
    try:
        due = int(context.args[0])
        if due < 0:
            await update.message.reply_text('Sorry we can not go back to future!')
            return

        if context.job_queue is None:
            await update.message.reply_text('Notifications not available: APScheduler not installed.')
            return
        for job in context.job_queue.get_jobs_by_name(str(chat_id)):
            job.schedule_removal()
        context.job_queue.run_repeating(alarm, interval=due, first=due,
                                        chat_id=chat_id, name=str(chat_id))

        api_url = get_api_url()
        token = get_api_token()
        if api_url and token:
            params = urllib.parse.urlencode({'limit': 5})
            d = fetch_json(f"{api_url}/social_log?{params}", token)
            if d:
                entries = d if isinstance(d, list) else d.get('entries', [])
                max_id = entries[-1].get('id', 0) if entries else 0
                social_log_last_ids[chat_id] = max_id

        await update.message.reply_text('Notifications enabled!')

    except (IndexError, ValueError):
        await update.message.reply_text('Usage: /set <seconds>')


async def unset_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):
    chat_id = update.message.chat_id
    jobs = context.job_queue.get_jobs_by_name(str(chat_id))
    if not jobs:
        await update.message.reply_text('Notifications were not enabled')
        return
    for job in jobs:
        job.schedule_removal()
    await update.message.reply_text('Notifications disabled')


async def ask(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Handle /ask <question>: explicit AI assistant invocation."""
    lang = get_chat_language(update.effective_chat.id)
    if not context.args:
        no_args = {'en': 'Usage: /ask <question>', 'it': 'Utilizzo: /ask <domanda>'}
        await update.message.reply_text(no_args[lang])
        return
    await _do_ask(update, context, ' '.join(context.args))


async def reset(update: Update, context: ContextTypes.DEFAULT_TYPE):
    """Clear the AI conversation history for this chat."""
    lang = get_chat_language(update.effective_chat.id)
    chat_id = update.effective_chat.id
    conversation_history.pop(chat_id, None)
    msg = {'en': 'Conversation history cleared.', 'it': 'Cronologia conversazione cancellata.'}
    await update.message.reply_text(msg[lang])


async def error_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
    logger.warning('Update "%s" caused error "%s"' % (update, context.error))


def main():
    config = get_config()
    bot_token = config.get('TELEGRAM_BOT_TOKEN', '')
    if not bot_token:
        logger.error("TELEGRAM_BOT_TOKEN not set in %s" % CONFIG_FILE)
        return

    application = Application.builder().token(bot_token).build()

    application.add_handler(CommandHandler('start', start))
    application.add_handler(CommandHandler('help', help_cmd))
    application.add_handler(CommandHandler('search', search_cmd))
    application.add_handler(CommandHandler('details', details))
    application.add_handler(CommandHandler('ask', ask))
    application.add_handler(CommandHandler('reset', reset))
    application.add_handler(CommandHandler('set', set_cmd))
    application.add_handler(CommandHandler('unset', unset_cmd))
    application.add_handler(MessageHandler(
        filters.TEXT & ~filters.COMMAND & filters.UpdateType.MESSAGES, query))

    application.add_error_handler(error_handler)

    application.run_polling()


if __name__ == '__main__':
    main()
