app.py

import discord, requests, os, aiosqlite, time, hmac, json
from discord.ext import commands
from discord import app_commands
from aiohttp import web

CEREBRAS_KEY = os.environ.get("CEREBRAS_KEY")
DISCORD_TOKEN = os.environ.get("DISCORD_TOKEN")
WEBHOOK_SECRET = os.environ.get("WEBHOOK_SECRET")

DAILY_FREE_TOKENS = 1000
MAX_TOKENS_PER_REQ = 250
CHAT_MAX_TOKENS = 200
TOKEN_PACK_AMOUNT = 1000 # ile daje 1x "Token AI Pack" z UB

intents = discord.Intents.default()
bot = commands.Bot(command_prefix="!", intents=intents)

# === BAZA DANYCH ===
async def init_db():
    async with aiosqlite.connect("data.db") as db:
        await db.execute('''CREATE TABLE IF NOT EXISTS users (
            user_id INTEGER PRIMARY KEY,
            tokens INTEGER DEFAULT 0,
            last_reset INTEGER DEFAULT 0
        )''')
        await db.execute('''CREATE TABLE IF NOT EXISTS chat_sessions (
            user_id INTEGER PRIMARY KEY,
            channel_id INTEGER,
            history TEXT DEFAULT '[]'
        )''')
        await db.execute('''CREATE TABLE IF NOT EXISTS used_webhooks (
            inventory_id TEXT PRIMARY KEY
        )''')
        await db.commit()

async def get_user_tokens(user_id: int):
    async with aiosqlite.connect("data.db") as db:
        async with db.execute("SELECT tokens, last_reset FROM users WHERE user_id =?", (user_id,)) as cur:
            row = await cur.fetchone()
        now = int(time.time())
        if not row or now - row[1] > 86400:
            await db.execute("INSERT OR REPLACE INTO users VALUES (?,?,?)",
                             (user_id, DAILY_FREE_TOKENS, now))
            await db.commit()
            return DAILY_FREE_TOKENS
        return row[0]

async def use_tokens(user_id: int, amount: int):
    async with aiosqlite.connect("data.db") as db:
        await db.execute("UPDATE users SET tokens = tokens -? WHERE user_id =?", (amount, user_id))
        await db.commit()

async def add_tokens(user_id: int, amount: int):
    async with aiosqlite.connect("data.db") as db:
        await db.execute("INSERT INTO users (user_id, tokens, last_reset) VALUES (?,?,?) ON CONFLICT(user_id) DO UPDATE SET tokens = tokens +?",
                         (user_id, amount, int(time.time()), amount))
        await db.commit()

# === KOMENDY ===
@bot.event
async def on_ready():
    await init_db()
    bot.loop.create_task(start_webserver())
    await bot.tree.sync()
    print(f'Bot online: {bot.user}')

@bot.tree.command(name="ask", description="Zapytaj AI. Limit: 1000 tokenów/dzień")
@app_commands.checks.cooldown(1, 20, key=lambda i: i.user.id)
@app_commands.describe(prompt="Twoje pytanie do AI")
async def ask(interaction: discord.Interaction, prompt: str):
    await interaction.response.defer()
    user_id = interaction.user.id
    tokens_left = await get_user_tokens(user_id)
    if tokens_left < 50:
        await interaction.followup.send(f"Masz tylko {tokens_left} tokenów. Doładuj w sklepie lub wróć jutro.")
        return
    try:
        r = requests.post("https://api.cerebras.ai/v1/chat/completions",
            headers={"Authorization": f"Bearer {CEREBRAS_KEY}"},
            json={
                "model":"llama3.1-70b",
                "messages":[{"role":"user","content":prompt}],
                "max_tokens": min(MAX_TOKENS_PER_REQ, tokens_left)
            }, timeout=30)
        r.raise_for_status()
        reply = r.json()['choices'][0]['message']['content']
        used = r.json()['usage']['completion_tokens']
        await use_tokens(user_id, used)
        embed = discord.Embed(description=reply[:4000], color=0x2b2d31)
        embed.set_footer(text=f"Zużyto {used} tokenów | Zostało: {tokens_left - used}/1000")
        await interaction.followup.send(embed=embed)
    except Exception as e:
        await interaction.followup.send(f"Błąd: {e}")

@ask.error
async def ask_error(interaction, error):
    if isinstance(error, commands.CommandOnCooldown):
        await interaction.response.send_message(f"Zwolnij. Spróbuj za {error.retry_after:.0f}s", ephemeral=True)

@bot.tree.command(name="balance", description="Sprawdź ile masz tokenów AI")
async def balance(interaction: discord.Interaction):
    tokens = await get_user_tokens(interaction.user.id)
    await interaction.response.send_message(f"Masz **{tokens}/1000** tokenów AI na dziś.", ephemeral=True)

# === TRYB CZATU ===
chat_group = app_commands.Group(name="chat", description="Tryb swobodnej rozmowy z AI")

@chat_group.command(name="start", description="Włącz tryb czatu w tym kanale")
async def chat_start(interaction: discord.Interaction):
    user_id = interaction.user.id
    if await get_user_tokens(user_id) < 50:
        await interaction.response.send_message("Masz za mało tokenów żeby zacząć czat. Minimum 50.", ephemeral=True)
        return
    async with aiosqlite.connect("data.db") as db:
        await db.execute("INSERT OR REPLACE INTO chat_sessions VALUES (?,?,'[]')",
                         (user_id, interaction.channel_id))
        await db.commit()
    await interaction.response.send_message("Tryb czatu ON ✅ Pisz normalnie. `/chat stop` żeby zakończyć.", ephemeral=True)

@chat_group.command(name="stop", description="Wyłącz tryb czatu")
async def chat_stop(interaction: discord.Interaction):
    async with aiosqlite.connect("data.db") as db:
        await db.execute("DELETE FROM chat_sessions WHERE user_id =?", (interaction.user.id,))
        await db.commit()
    await interaction.response.send_message("Tryb czatu OFF.", ephemeral=True)

bot.tree.add_command(chat_group)

@bot.event
async def on_message(message: discord.Message):
    if message.author.bot: return
    async with aiosqlite.connect("data.db") as db:
        async with db.execute("SELECT channel_id, history FROM chat_sessions WHERE user_id =?",
                              (message.author.id,)) as cur:
            row = await cur.fetchone()
    if not row or row[0]!= message.channel.id:
        await bot.process_commands(message)
        return

    user_id = message.author.id
    tokens_left = await get_user_tokens(user_id)
    if tokens_left < 30:
        await message.reply("Skończyły Ci się tokeny. `/chat stop` i doładuj w sklepie.")
        return

    history = json.loads(row[1])
    history.append({"role": "user", "content": message.content})
    history = history[-6:] # pamięć 3 wymian

    async with message.channel.typing():
        try:
            r = requests.post("https://api.cerebras.ai/v1/chat/completions",
                headers={"Authorization": f"Bearer {CEREBRAS_KEY}"},
                json={
                    "model":"llama3.1-70b",
                    "messages": [{"role":"system","content":"Jesteś luzackim botem Discord. Odpowiadaj krótko, na temat, bez lania wody."}] + history,
                    "max_tokens": min(CHAT_MAX_TOKENS, tokens_left)
                }, timeout=30)
            r.raise_for_status()
            reply = r.json()['choices'][0]['message']['content']
            used = r.json()['usage']['completion_tokens']
            await use_tokens(user_id, used)
            history.append({"role": "assistant", "content": reply})
            async with aiosqlite.connect("data.db") as db:
                await db.execute("UPDATE chat_sessions SET history =? WHERE user_id =?",
                                 (json.dumps(history), user_id))
                await db.commit()
            await message.reply(reply[:2000])
        except Exception as e:
            await message.reply(f"AI padło: {e}")

# === WEBHOOK UNBELIEVABOAT ===
async def ub_webhook(request):
    secret = request.rel_url.query.get('key')
    if not WEBHOOK_SECRET or not hmac.compare_digest(secret, WEBHOOK_SECRET):
        return web.Response(text="Unauthorized", status=401)
    try:
        data = await request.json()
    except:
        return web.Response(text="Bad JSON", status=400)
    user_id = int(data.get('user_id', 0))
    item_name = data.get('item_name', '')
    quantity = int(data.get('quantity', 1))
    inventory_id = data.get('id')
    if not user_id or "Token AI" not in item_name:
        return web.Response(text="Ignored", status=200)
    async with aiosqlite.connect("data.db") as db:
        async with db.execute("SELECT 1 FROM used_webhooks WHERE inventory_id =?", (inventory_id,)) as cur:
            if await cur.fetchone():
                return web.Response(text="Already processed", status=200)
        await db.execute("INSERT INTO used_webhooks VALUES (?)", (inventory_id,))
        await db.commit()
    tokens_to_add = TOKEN_PACK_AMOUNT * quantity
    await add_tokens(user_id, tokens_to_add)
    return web.Response(text=f"Added {tokens_to_add} tokens")

async def start_webserver():
    app = web.Application()
    app.router.add_post('/ub-webhook', ub_webhook)
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, '0.0.0.0', 7860)
    await site.start()

bot.run(DISCORD_TOKEN)