"""Klarbildbot – Telegram-Bot fuer AI-Bildaufbereitung (Freistellen + Upscaling)."""
from __future__ import annotations

import asyncio
import logging
import posixpath
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

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

from . import ai_edit, config, enhance, picdrop

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

# Schwere CPU-Jobs strikt serialisieren (Speicherschutz auf kleinen Boxen).
EXECUTOR = ThreadPoolExecutor(max_workers=1)

# Kurzzeit-Speicher: pending-Bild pro User (file_id + Herkunft).
PENDING: dict[int, dict] = {}


# ---------------------------------------------------------------------------
# Zugriffsschutz
# ---------------------------------------------------------------------------
def _allowed(update: Update) -> bool:
    if not config.ALLOWED_USER_IDS:
        return True
    user = update.effective_user
    return bool(user and user.id in config.ALLOWED_USER_IDS)


async def _deny(update: Update) -> None:
    uid = update.effective_user.id if update.effective_user else "?"
    await update.effective_message.reply_text(
        f"⛔️ Kein Zugriff. Deine Telegram-ID: {uid}\n"
        "Bitte vom Betreiber in ALLOWED_USER_IDS freischalten lassen."
    )


# ---------------------------------------------------------------------------
# Commands
# ---------------------------------------------------------------------------
async def cmd_start(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
    uid = update.effective_user.id if update.effective_user else "?"
    if not _allowed(update):
        return await _deny(update)
    ai_block = ""
    if config.AI_EDIT_ENABLED:
        ai_block = (
            "• 🎨 *Studio-Hintergrund* / 🪄 *AI-Retusche* / 🖼 *Restaurierung*\n"
            "• 💬 *Eigener AI-Prompt*: Bild mit Bildunterschrift senden "
            "(z. B. 'Hintergrund zu Strand', 'mach es schwarz-weiß')\n"
        )
    await update.effective_message.reply_text(
        "👋 *Klarbildbot* – AI-Bildaufbereitung\n\n"
        "Schick mir ein Bild, dann kannst du wählen:\n"
        "• 🔍 *Klarbild ×2 / ×4* – Upscaling + Schärfen (lokal, originaltreu)\n"
        "• ✂️ *Freistellen* – Hintergrund entfernen, echtes transparentes PNG (lokal)\n"
        "• ✨ *Freistellen + Klarbild* – beides\n"
        + ai_block +
        "\n💡 Für beste Qualität das Bild als *Datei* senden "
        "(Büroklammer → Datei), nicht als komprimiertes Foto.\n\n"
        f"Deine Telegram-ID: `{uid}`",
        parse_mode="Markdown",
    )


async def cmd_help(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
    if not _allowed(update):
        return await _deny(update)
    txt = (
        "*So geht's:*\n"
        "1. Bild senden (am besten als Datei)\n"
        "2. Aktion wählen\n"
        "3. Ergebnis kommt als Datei zurück\n\n"
        "*Befehle:*\n"
        "/start – Übersicht\n"
        "/help – diese Hilfe\n"
    )
    if config.PICDROP_ENABLED:
        txt += "/picdrop – Picdrop-Galerie stapelweise verarbeiten\n"
    await update.effective_message.reply_text(txt, parse_mode="Markdown")


# ---------------------------------------------------------------------------
# Bild empfangen -> Aktionsmenue
# ---------------------------------------------------------------------------
def _menu() -> InlineKeyboardMarkup:
    rows = [
        [InlineKeyboardButton("🔍 Klarbild ×2", callback_data="up:2"),
         InlineKeyboardButton("🔍 Klarbild ×4", callback_data="up:4")],
        [InlineKeyboardButton("✂️ Freistellen", callback_data="cut:0")],
        [InlineKeyboardButton("✨ Freistellen + Klarbild ×2", callback_data="cutup:2")],
    ]
    if config.AI_EDIT_ENABLED:
        rows += [
            [InlineKeyboardButton("🎨 Studio-Hintergrund", callback_data="ai:studio"),
             InlineKeyboardButton("🪄 AI-Retusche", callback_data="ai:retouch")],
            [InlineKeyboardButton("🖼 Foto-Restaurierung", callback_data="ai:restore")],
        ]
    return InlineKeyboardMarkup(rows)


def _label(action: str, scale: int = 0, prompt: str = "") -> str:
    return {
        "up": f"Klarbild ×{scale}",
        "cut": "Freistellen",
        "cutup": f"Freistellen + Klarbild ×{scale}",
        "ai": ai_edit.PRESET_LABELS.get(prompt, prompt),
        "aiprompt": f"AI-Edit: {prompt[:40]}",
    }.get(action, action)


def _do_job(action: str, scale: int, prompt: str, image_bytes: bytes) -> enhance.Result:
    if action == "up":
        return enhance.upscale(image_bytes, scale)
    if action == "cut":
        return enhance.cutout(image_bytes)
    if action == "cutup":
        return enhance.cutout_then_upscale(image_bytes, scale)
    if action == "ai":
        return ai_edit.edit(image_bytes, ai_edit.PRESETS[prompt])
    if action == "aiprompt":
        return ai_edit.edit(image_bytes, prompt)
    raise ValueError(action)


async def _process(ctx: ContextTypes.DEFAULT_TYPE, chat_id: int, status,
                   file_id: str, action: str, scale: int, prompt: str) -> None:
    """Bild laden, Job im Executor ausfuehren, Ergebnis als Dokument senden.
    `status` ist eine Telegram-Message, deren Text als Fortschritt editiert wird."""
    label = _label(action, scale, prompt)
    await status.edit_text(f"⏳ {label} läuft … (kann bei großen Bildern dauern)")
    try:
        tg_file = await ctx.bot.get_file(file_id)
        image_bytes = bytes(await tg_file.download_as_bytearray())
    except Exception as exc:
        log.exception("Download fehlgeschlagen")
        return await status.edit_text(
            f"❌ Konnte das Bild nicht laden: {exc}\n"
            "(Telegram-Bots können Dateien bis 20 MB laden.)")

    await ctx.bot.send_chat_action(chat_id, ChatAction.UPLOAD_DOCUMENT)
    loop = asyncio.get_running_loop()
    try:
        result = await loop.run_in_executor(
            EXECUTOR, _do_job, action, scale, prompt, image_bytes)
    except Exception as exc:
        log.exception("Verarbeitung fehlgeschlagen")
        return await status.edit_text(f"❌ Fehler bei der Verarbeitung: {exc}")

    caption = (f"✅ {label}\n{result.width}×{result.height}px · {result.seconds:.1f}s"
               + (f"\n{result.note}" if result.note else ""))
    try:
        await ctx.bot.send_document(
            chat_id=chat_id,
            document=InputFile(result.data, filename=f"klarbild.{result.fmt}"),
            caption=caption,
        )
        await status.edit_text(f"✅ Fertig: {label}")
    except Exception as exc:
        log.exception("Senden fehlgeschlagen")
        await status.edit_text(f"❌ Konnte Ergebnis nicht senden: {exc}")


async def on_image(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
    if not _allowed(update):
        return await _deny(update)
    msg = update.effective_message

    if msg.photo:
        file_id = msg.photo[-1].file_id
        source = "photo"
    elif msg.document and (msg.document.mime_type or "").startswith("image/"):
        file_id = msg.document.file_id
        source = "document"
    else:
        return await msg.reply_text("Bitte ein Bild senden.")

    # Bild mit Bildunterschrift + AI aktiv => Unterschrift als freier AI-Prompt
    caption = (msg.caption or "").strip()
    if caption and config.AI_EDIT_ENABLED:
        status = await msg.reply_text("🪄 AI-Edit vorbereiten …")
        return await _process(ctx, msg.chat_id, status, file_id, "aiprompt", 0, caption)

    PENDING[update.effective_user.id] = {"file_id": file_id, "source": source}
    hint = "" if source == "document" else "\n_(Tipp: als Datei senden = volle Auflösung)_"
    if config.AI_EDIT_ENABLED:
        hint += "\n💬 _Oder Bild mit Bildunterschrift senden = eigener AI-Prompt._"
    await msg.reply_text(
        "Was soll ich damit machen?" + hint,
        reply_markup=_menu(),
        parse_mode="Markdown",
    )


# ---------------------------------------------------------------------------
# Button gedrueckt -> verarbeiten
# ---------------------------------------------------------------------------
async def on_action(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
    query = update.callback_query
    await query.answer()
    if not _allowed(update):
        return await query.edit_message_text("⛔️ Kein Zugriff.")

    uid = update.effective_user.id
    pend = PENDING.get(uid)
    if not pend:
        return await query.edit_message_text(
            "Kein Bild gemerkt – bitte schick es nochmal.")

    action, _, arg = query.data.partition(":")
    scale = int(arg) if arg.isdigit() else 0
    prompt = arg if action == "ai" else ""   # bei AI-Presets ist arg der Preset-Key

    try:
        await _process(ctx, update.effective_chat.id, query.message,
                       pend["file_id"], action, scale, prompt)
    finally:
        PENDING.pop(uid, None)


# ---------------------------------------------------------------------------
# Picdrop-Batch
# ---------------------------------------------------------------------------
async def cmd_picdrop(update: Update, ctx: ContextTypes.DEFAULT_TYPE) -> None:
    if not _allowed(update):
        return await _deny(update)
    if not config.PICDROP_ENABLED:
        return await update.effective_message.reply_text(
            "Picdrop ist nicht konfiguriert (PICDROP_USER / PICDROP_PASSWORD).")

    args = ctx.args or []
    loop = asyncio.get_running_loop()

    if not args:
        try:
            entries = await loop.run_in_executor(EXECUTOR, picdrop.list_dir, "/")
        except Exception as exc:
            return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
        folders = [e["name"] for e in entries if e["is_dir"]][:40]
        listing = "\n".join(f"• {f}" for f in folders) or "(keine Ordner)"
        return await update.effective_message.reply_text(
            "*Picdrop – Ordner im Wurzelverzeichnis:*\n" + listing +
            "\n\nNutzung: `/picdrop <Ordnername-Teil> up2|up4|cut`",
            parse_mode="Markdown")

    needle = args[0]
    op = (args[1] if len(args) > 1 else "up2").lower()

    await update.effective_message.reply_text(f"🔎 Suche Galerie „{needle}“ …")
    try:
        matches = await loop.run_in_executor(EXECUTOR, picdrop.find_gallery, needle)
    except Exception as exc:
        return await update.effective_message.reply_text(f"❌ SFTP-Fehler: {exc}")
    if not matches:
        return await update.effective_message.reply_text("Keine passende Galerie gefunden.")

    gallery = matches[0]
    images = await loop.run_in_executor(EXECUTOR, picdrop.list_images, gallery)
    if not images:
        return await update.effective_message.reply_text(
            f"Galerie `{gallery}` enthält keine Bilder.", parse_mode="Markdown")

    await update.effective_message.reply_text(
        f"📁 `{gallery}` · {len(images)} Bilder · Aktion `{op}`\n"
        "Verarbeite und lade in Unterordner `klarbild/` hoch …",
        parse_mode="Markdown")

    def process_one(remote_path: str) -> str:
        data = picdrop.download(remote_path)
        if op == "cut":
            res = enhance.cutout(data)
        elif op == "up4":
            res = enhance.upscale(data, 4)
        else:
            res = enhance.upscale(data, 2)
        base = posixpath.splitext(posixpath.basename(remote_path))[0]
        out_name = f"{base}_klar.{res.fmt}"
        out_path = posixpath.join(gallery, "klarbild", out_name)
        picdrop.upload(out_path, res.data)
        return out_name

    done = 0
    for img in images:
        try:
            name = await loop.run_in_executor(EXECUTOR, process_one, img["path"])
            done += 1
            if done % 5 == 0 or done == len(images):
                await update.effective_message.reply_text(
                    f"… {done}/{len(images)} fertig (zuletzt: {name})")
        except Exception as exc:
            log.exception("Picdrop-Bild fehlgeschlagen")
            await update.effective_message.reply_text(
                f"⚠️ {img['name']}: {exc}")

    await update.effective_message.reply_text(
        f"✅ Fertig: {done}/{len(images)} Bilder → `{gallery}/klarbild/`",
        parse_mode="Markdown")


# ---------------------------------------------------------------------------
# Health-Server (fuer Coolify) + App-Setup
# ---------------------------------------------------------------------------
async def _start_health(app: Application) -> None:
    from aiohttp import web

    async def health(_req):
        return web.json_response({"status": "ok", "service": "klarbildbot"})

    server = web.Application()
    server.router.add_get("/", health)
    server.router.add_get("/health", health)
    runner = web.AppRunner(server)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", config.HEALTH_PORT)
    await site.start()
    app.bot_data["health_runner"] = runner
    log.info("Health-Server auf :%s", config.HEALTH_PORT)


async def _post_init(app: Application) -> None:
    await _start_health(app)
    # Bot-Kommandos in Telegram-UI registrieren
    from telegram import BotCommand

    cmds = [BotCommand("start", "Übersicht"), BotCommand("help", "Hilfe")]
    if config.PICDROP_ENABLED:
        cmds.append(BotCommand("picdrop", "Picdrop-Galerie verarbeiten"))
    await app.bot.set_my_commands(cmds)
    # Modelle vorwaermen (im Thread, blockiert den Loop nicht)
    asyncio.get_running_loop().run_in_executor(EXECUTOR, enhance.warmup)
    log.info("Klarbildbot bereit.")


async def _post_shutdown(app: Application) -> None:
    runner = app.bot_data.get("health_runner")
    if runner:
        await runner.cleanup()


def build_app() -> Application:
    config.validate()
    app = (
        Application.builder()
        .token(config.BOT_TOKEN)
        .post_init(_post_init)
        .post_shutdown(_post_shutdown)
        .build()
    )
    app.add_handler(CommandHandler("start", cmd_start))
    app.add_handler(CommandHandler("help", cmd_help))
    app.add_handler(CommandHandler("picdrop", cmd_picdrop))
    app.add_handler(MessageHandler(
        filters.PHOTO | filters.Document.IMAGE, on_image))
    app.add_handler(CallbackQueryHandler(on_action))
    return app


def main() -> None:
    app = build_app()
    log.info("Starte Polling …")
    app.run_polling(allowed_updates=Update.ALL_TYPES, drop_pending_updates=True)


if __name__ == "__main__":
    main()