import os
import tempfile
import shutil
import threading
import time
import requests

from flask import Flask

from pyrogram import Client, filters
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton

from huggingface_hub import HfApi, hf_hub_download

# =========================
# CONFIG
# =========================

API_ID = int(os.environ["API_ID"])
API_HASH = os.environ["API_HASH"]
BOT_TOKEN = os.environ["BOT_TOKEN"]
HF_TOKEN = os.environ["HF_TOKEN"]

DATASET_REPO = "Cxops/Bzxv"

# ✅ FIXED: correct Space runtime URL (IMPORTANT)
KEEP_ALIVE_URL = "https://cxops-push.hf.space/"
PING_INTERVAL = 300  # 5 minutes

# =========================
# FLASK HEALTH SERVER
# =========================

web = Flask(__name__)

@web.route("/")
def home():
    return "Bot Running..."

def run_web():
    web.run(host="0.0.0.0", port=7860)

# =========================
# KEEP ALIVE PINGER
# =========================

def keep_alive():
    headers = {
        "User-Agent": "Mozilla/5.0 (KeepAliveBot)"
    }

    while True:
        try:
            r = requests.get(KEEP_ALIVE_URL, headers=headers, timeout=20)
            print(f"[KEEPALIVE] Pinged {KEEP_ALIVE_URL} | Status: {r.status_code}")
        except Exception as e:
            print(f"[KEEPALIVE ERROR] {e}")

        time.sleep(PING_INTERVAL)

# =========================
# HUGGINGFACE
# =========================

hf = HfApi(token=HF_TOKEN)

# =========================
# PYROGRAM BOT
# =========================

app = Client(
    "storage_bot",
    api_id=API_ID,
    api_hash=API_HASH,
    bot_token=BOT_TOKEN
)

# =========================
# CLEANUP HELPER
# =========================

def cleanup(path):
    try:
        if os.path.isdir(path):
            shutil.rmtree(path, ignore_errors=True)
        elif os.path.isfile(path):
            os.remove(path)
    except Exception:
        pass

# =========================
# COMMANDS
# =========================

@app.on_message(filters.command("start"))
async def start(_, message):
    await message.reply_text(
        "Send me a file.\n\n"
        "/files - List uploaded files"
    )

@app.on_message(filters.command("files"))
async def files(_, message):
    user_id = str(message.from_user.id)

    try:
        repo_files = hf.list_repo_files(
            repo_id=DATASET_REPO,
            repo_type="dataset"
        )

        user_files = [
            f for f in repo_files
            if f.startswith(f"{user_id}/")
        ]

        if not user_files:
            return await message.reply_text("You have no uploaded files.")

        buttons = [
            [
                InlineKeyboardButton(
                    os.path.basename(path)[:60],
                    callback_data=f"get:{path}"
                )
            ]
            for path in user_files
        ]

        await message.reply_text(
            "Your uploaded files:",
            reply_markup=InlineKeyboardMarkup(buttons)
        )

    except Exception as e:
        await message.reply_text(f"Error:\n{e}")

# =========================
# DOWNLOAD FROM HF
# =========================

@app.on_callback_query(filters.regex("^get:"))
async def get_file(_, query):
    path = query.data.replace("get:", "")

    temp_dir = tempfile.mkdtemp()

    try:
        await query.message.edit_text("Downloading from dataset...")

        local_file = hf_hub_download(
            repo_id=DATASET_REPO,
            repo_type="dataset",
            filename=path,
            token=HF_TOKEN,
            local_dir=temp_dir
        )

        await query.message.reply_document(
            local_file,
            caption=os.path.basename(path)
        )

        try:
            await query.message.delete()
        except Exception:
            pass

    except Exception as e:
        await query.message.reply_text(f"Error:\n{e}")

    finally:
        cleanup(temp_dir)

# =========================
# FILE UPLOAD HANDLER
# =========================

@app.on_message(
    filters.document |
    filters.video |
    filters.audio |
    filters.photo
)
async def upload_file(_, message):

    temp_dir = tempfile.mkdtemp()

    try:
        status = await message.reply_text("Downloading file...")

        if message.document:
            media = message.document
            filename = media.file_name

        elif message.video:
            media = message.video
            filename = media.file_name or f"{media.file_unique_id}.mp4"

        elif message.audio:
            media = message.audio
            filename = media.file_name or f"{media.file_unique_id}.mp3"

        else:
            media = message.photo[-1]
            filename = f"{media.file_unique_id}.jpg"

        local_path = await message.download(
            file_name=os.path.join(temp_dir, filename)
        )

        await status.edit_text("Uploading to Hugging Face...")

        repo_path = f"{message.from_user.id}/{filename}"

        hf.upload_file(
            path_or_fileobj=local_path,
            path_in_repo=repo_path,
            repo_id=DATASET_REPO,
            repo_type="dataset"
        )

        await status.edit_text(f"Uploaded successfully.\n\n{filename}")

    except Exception as e:
        await message.reply_text(f"Upload failed:\n{e}")

    finally:
        cleanup(temp_dir)

# =========================
# MAIN
# =========================

if __name__ == "__main__":
    print("Starting Flask health server...")

    threading.Thread(target=run_web, daemon=True).start()

    print("Starting Keep-Alive Pinger...")

    threading.Thread(target=keep_alive, daemon=True).start()

    print("Starting Telegram bot...")

    app.run()