import os
import re
import html
import time
import threading
import requests

from flask import Flask
from pyrogram import Client, filters
from pyrogram.enums import ParseMode
from pyrogram.types import InlineKeyboardMarkup, InlineKeyboardButton

# ======================
# ENV VARIABLES
# ======================

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

CVE_API_BASE = os.environ["CVE_API_BASE"]

# ======================
# HF SPACE URL (FILLED)
# ======================

SPACE_URL = "https://cistern-opa.hf.space/"

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

web = Flask(__name__)

@web.route("/")
def home():
    return {"status": "running"}

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

# ======================
# HEARTBEAT (REAL KEEP ALIVE)
# ======================

def heartbeat():
    while True:
        try:
            r = requests.get(SPACE_URL, timeout=10)
            print(f"Keep-alive ping sent: {r.status_code}")
        except Exception as e:
            print(f"Keep-alive failed: {e}")

        time.sleep(300)

# ======================
# BOT
# ======================

app = Client(
    "cvebot",
    api_id=API_ID,
    api_hash=API_HASH,
    bot_token=BOT_TOKEN,
    parse_mode=ParseMode.HTML
)

# ======================
# HELPERS
# ======================

def sev_emoji(sev):
    sev = str(sev).upper()
    return {
        "CRITICAL": "🔴",
        "HIGH": "🟠",
        "MEDIUM": "🟡",
        "LOW": "🟢"
    }.get(sev, "⚪")


def get_desc(cna):
    for d in cna.get("descriptions", []):
        if d.get("lang") == "en":
            return d.get("value", "")
    return "No description available."


def get_cvss(cna):
    for m in cna.get("metrics", []):
        for _, v in m.items():
            if isinstance(v, dict) and "baseScore" in v:
                return v
    return None


def get_cwes(cna):
    out = []
    for p in cna.get("problemTypes", []):
        for d in p.get("descriptions", []):
            if d.get("cweId"):
                out.append(d["cweId"])
    return out


def get_products(cna):
    out = []
    for a in cna.get("affected", []):
        out.append(f"{a.get('vendor','')} {a.get('product','')}".strip())
    return out


def get_refs(cna):
    return [r["url"] for r in cna.get("references", []) if r.get("url")]

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

@app.on_message(filters.command("start"))
async def start(_, msg):
    await msg.reply_text(
        "🔎 CVE Bot Ready\n\nUse:\n/cve CVE-2021-1256"
    )


@app.on_message(filters.command("cve"))
async def cve_lookup(_, msg):

    if len(msg.command) < 2:
        return await msg.reply_text("Usage: /cve CVE-XXXX-XXXX")

    cve_id = msg.command[1].upper()

    if not re.match(r"CVE-\d{4}-\d+", cve_id):
        return await msg.reply_text("❌ Invalid CVE format")

    url = f"{CVE_API_BASE}{cve_id}"

    try:
        r = requests.get(url, timeout=20)

        if r.status_code != 200:
            return await msg.reply_text(f"❌ CVE not found ({r.status_code})")

        data = r.json()

        cna = data.get("containers", {}).get("cna", {})
        meta = data.get("cveMetadata", {})

        desc = html.escape(get_desc(cna))
        cwes = get_cwes(cna)
        products = get_products(cna)
        refs = get_refs(cna)

        cvss = get_cvss(cna)

        text = f"<b>🚨 {cve_id}</b>\n\n"

        if cna.get("title"):
            text += f"<b>📌 {html.escape(cna['title'])}</b>\n\n"

        if cvss:
            sev = cvss.get("baseSeverity", "Unknown")
            score = cvss.get("baseScore", "N/A")

            text += (
                f"<b>🎯 Severity</b>\n"
                f"{sev_emoji(sev)} {sev}\n"
                f"Score: <code>{score}</code>\n\n"
            )

        if cwes:
            text += "<b>⚠️ CWE</b>\n"
            text += "\n".join(f"• {c}" for c in cwes[:10]) + "\n\n"

        if products:
            text += "<b>📦 Affected Products</b>\n"
            text += "\n".join(f"• {p}" for p in products[:10]) + "\n\n"

        text += (
            f"<b>📝 Description</b>\n"
            f"<blockquote>{desc[:3000]}</blockquote>\n\n"
            f"<b>📅 Published</b>\n<code>{meta.get('datePublished','N/A')}</code>\n\n"
            f"<b>🔄 Updated</b>\n<code>{meta.get('dateUpdated','N/A')}</code>\n"
        )

        buttons = [[InlineKeyboardButton("🔗 MITRE CVE", url=url)]]

        if refs:
            buttons.append([InlineKeyboardButton("📎 Reference", url=refs[0])])

        await msg.reply_text(
            text,
            reply_markup=InlineKeyboardMarkup(buttons)
        )

    except Exception as e:
        await msg.reply_text(f"❌ Error: {e}")

# ======================
# START
# ======================

if __name__ == "__main__":

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

    print("Bot Started")
    app.run()