import aiohttp
import logging

from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
from telegram.constants import ChatAction
from telegram.ext import (
    Application,
    CommandHandler,
    MessageHandler,
    CallbackQueryHandler,
    ContextTypes,
    filters,
)


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s"
)


BOT_TOKEN = "8962551714:AAGlmUXvyERWBrmq1wgtB0U21veWXs76pGA"

API_URL = "https://py.swifthostbd.store/v1/chat/completions"


user_model = {}


MODELS = {
    "gpt55": "gpt-5.5",
    "gpt54": "gpt-5.4",
    "gpt54m": "gpt-5.4-mini",
    "gpt54n": "gpt-5.4-nano",

    "gem35": "gemini-3.5-flash",
    "gem3": "gemini-3-flash",
    "gem31": "gemini-3.1-flash-lite",
    "gem25": "gemini-2.5-flash",
    "gem25l": "gemini-2.5-flash-lite",

    "gemma26": "gemma-4-26b",
    "gemma31": "gemma-4-31b",

    "olmo31": "Olmo-3.1-32B-Instruct",
    "olmo7": "Olmo-3-7B-Instruct",
    "molmo": "molmo2-8b",

    "trmini": "trinity-mini",
    "trthink": "trinity-large-thinking",

    "nem30": "nemotron-3-nano-30b",
    "nem120": "nemotron-3-super-120b",

    "sar30": "sarvam-30b",
    "sar105": "sarvam-105b",
}


def keyboard():

    return InlineKeyboardMarkup([

        [
            InlineKeyboardButton("🔥 GPT-5.5", callback_data="gpt55"),
            InlineKeyboardButton("⚡ GPT-5.4", callback_data="gpt54")
        ],

        [
            InlineKeyboardButton("🚀 GPT-5.4 Mini", callback_data="gpt54m"),
            InlineKeyboardButton("🟢 GPT-5.4 Nano", callback_data="gpt54n")
        ],

        [
            InlineKeyboardButton("💎 Gemini 3.5", callback_data="gem35"),
            InlineKeyboardButton("✨ Gemini 3", callback_data="gem3")
        ],

        [
            InlineKeyboardButton("🌙 Gemini 2.5", callback_data="gem25"),
            InlineKeyboardButton("🌙 Gemini Lite", callback_data="gem25l")
        ],

        [
            InlineKeyboardButton("💠 Gemma 31B", callback_data="gemma31"),
            InlineKeyboardButton("💠 Gemma 26B", callback_data="gemma26")
        ],

        [
            InlineKeyboardButton("🧠 Trinity Think", callback_data="trthink"),
            InlineKeyboardButton("⚡ Trinity Mini", callback_data="trmini")
        ],

        [
            InlineKeyboardButton("🦙 OLMo 32B", callback_data="olmo31"),
            InlineKeyboardButton("🦙 OLMo 7B", callback_data="olmo7")
        ],

        [
            InlineKeyboardButton("👁 Molmo2", callback_data="molmo")
        ],

        [
            InlineKeyboardButton("🟦 Nemotron 120B", callback_data="nem120"),
            InlineKeyboardButton("🟦 Nemotron 30B", callback_data="nem30")
        ],

        [
            InlineKeyboardButton("🌐 Sarvam 105B", callback_data="sar105"),
            InlineKeyboardButton("🌐 Sarvam 30B", callback_data="sar30")
        ]

    ])



async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):

    user_model[update.effective_user.id] = "gpt-5.5"

    await update.message.reply_text(
        "🤖 AI Model Selector\n\nChoose your model:",
        reply_markup=keyboard()
    )



async def help_cmd(update: Update, context: ContextTypes.DEFAULT_TYPE):

    await update.message.reply_text(
        """
Commands:

/start - Open model menu
/model - Change model
/help - Show help

Send any message to chat with AI.
"""
    )



async def model(update: Update, context: ContextTypes.DEFAULT_TYPE):

    await update.message.reply_text(
        "Select Model:",
        reply_markup=keyboard()
    )



async def button(update: Update, context: ContextTypes.DEFAULT_TYPE):

    query = update.callback_query

    await query.answer()


    selected = MODELS.get(
        query.data,
        "gpt-5.5"
    )


    user_model[query.from_user.id] = selected


    await query.edit_message_text(
        f"✅ Current Model:\n{selected}"
    )



async def chat(update: Update, context: ContextTypes.DEFAULT_TYPE):

    user_id = update.effective_user.id


    selected_model = user_model.get(
        user_id,
        "gpt-5.5"
    )


    await context.bot.send_chat_action(
        update.effective_chat.id,
        ChatAction.TYPING
    )


    payload = {

        "model": selected_model,

        "messages":[
            {
                "role":"user",
                "content":update.message.text
            }
        ]

    }


    try:

        timeout = aiohttp.ClientTimeout(total=60)


        async with aiohttp.ClientSession(timeout=timeout) as session:

            async with session.post(
                API_URL,
                json=payload
            ) as r:


                text = await r.text()


                try:
                    data = await r.json()
                except:
                    data = {}


                if "choices" in data:

                    reply = (
                        data["choices"][0]
                        ["message"]
                        ["content"]
                    )

                elif "error" in data:

                    reply = str(data["error"])

                else:

                    reply = text



    except Exception as e:

        reply = f"❌ Error:\n{e}"



    if len(reply) > 4000:
        reply = reply[:4000]


    await update.message.reply_text(reply)



async def error_handler(update, context):

    logging.error(
        "Exception: %s",
        context.error
    )



app = Application.builder().token(BOT_TOKEN).build()


app.add_handler(CommandHandler("start", start))
app.add_handler(CommandHandler("model", model))
app.add_handler(CommandHandler("help", help_cmd))

app.add_handler(
    CallbackQueryHandler(button)
)


app.add_handler(
    MessageHandler(
        filters.TEXT & ~filters.COMMAND,
        chat
    )
)


app.add_error_handler(error_handler)


print("🤖 Bot Started")

app.run_polling()