# latex_commands.py
# -*- coding: utf-8 -*-
from __future__ import annotations

import io
import logging
from typing import Optional

import discord
from discord import app_commands
from discord.ext import commands

from constants import (
    ALLOWED_ROLES,
    GUILD_ID,
    LATEX_MAX_ATTACH_MB,
    LATEX_MAX_COMPILE_SECONDS,
    LATEX_MAX_ZIP_MB,
    LATEX_TEX_ENGINE,
    OPENAI_MODEL,
)
from latex_functions import (
    compile_single_tex_bytes,
    compile_zip_to_zip,
    is_safe_asset_name,
    is_safe_tex_name,
)


async def _collect_attachments_bytes(attachments, max_mb: int):
    """
    Zwraca słownik {filename: bytes} dla wszystkich załączników, do limitu MB.
    Nie podbija wyjątków z .read() — zostawiamy to wyżej; tu zwracamy, co się udało.
    """
    out = {}
    for a in attachments or []:
        # limit rozmiaru
        if (getattr(a, "size", 0) or 0) > max_mb * 1024 * 1024:
            continue
        data = await a.read()
        out[a.filename] = data
    return out


class LatexModule(commands.Cog):
    def __init__(self, bot: commands.Bot, logger_name: str):
        self.bot = bot
        self.logger = logging.getLogger(logger_name)
        self.logger_name = logger_name

    # -------- /tex (PDF + diagnoza; log -> debug logger) --------
    @commands.hybrid_command(
        nsfw=False,
        name="tex",
        description="Kompiluje załączone .tex; PDF-y odsyła. Błędy: diagnoza z OpenAI.",
        guild=None if GUILD_ID is None else discord.Object(id=GUILD_ID),
    )
    @commands.has_any_role(*ALLOWED_ROLES)
    async def tex(
        self, ctx: commands.Context, include_source_excerpt: Optional[bool] = True
    ):
        self.logger.info("LaTeX command invoked by %s", ctx.author)
        ch = ctx.message.channel
        # ZBIERZ WSZYSTKIE ZAŁĄCZNIKI
        async with ch.typing():
            all_bytes = await _collect_attachments_bytes(
                ctx.message.attachments, LATEX_MAX_ATTACH_MB
            )
            if not all_bytes:
                return await ctx.reply(
                    f"Dołącz .tex (≤{LATEX_MAX_ATTACH_MB} MiB).", mention_author=False
                )

            # PODZIEL NA GŁÓWNE .TEX i WSPIERAJĄCE
            tex_names = [
                n
                for n in all_bytes.keys()
                if n.lower().endswith(".tex") and is_safe_tex_name(n)
            ]
            if not tex_names:
                return await ctx.reply(
                    f"Dołącz .tex (≤{LATEX_MAX_ATTACH_MB} MiB).", mention_author=False
                )

            support_names = [
                n
                for n in all_bytes.keys()
                if n not in tex_names and is_safe_asset_name(n)
            ]

            files_to_send = []
            parts = []

            for tex_name in tex_names:
                try:
                    main_bytes = all_bytes[tex_name]
                    support = [
                        (n, all_bytes[n])
                        for n in support_names
                        if n != tex_name and is_safe_asset_name(n)
                    ]
                    # (ważne) do supportów dopuszczamy również **inne .tex** – dla \input/\include
                    other_tex = [
                        (n, all_bytes[n])
                        for n in tex_names
                        if n != tex_name and is_safe_tex_name(n)
                    ]
                    support.extend(other_tex)

                    res = await compile_single_tex_bytes(
                        main_bytes,
                        tex_name,
                        LATEX_TEX_ENGINE,
                        LATEX_MAX_COMPILE_SECONDS,
                        logger="discord",
                        support_files=support,
                    )

                    # Twoja dotychczasowa obsługa success/fail – przykładowo:
                    if res["ok"] and res["pdf_bytes"]:
                        b = io.BytesIO(res["pdf_bytes"])
                        b.seek(0)
                        b.name = res["pdf_name"]
                        files_to_send.append(discord.File(b, filename=res["pdf_name"]))
                        parts.append(f"✅ `{tex_name}` → `{res['pdf_name']}`")
                    else:
                        parts.append(
                            f"❌ `{tex_name}` — błąd kompilacji.)"
                        )
                        # Tutaj wywołujesz swoją ścieżkę do OpenAI diagnozy, którą ogarnęliśmy wcześniej.
                except Exception:
                    # let it crash — pełny stack w logu i re-raise (Discord.py to pokaże sensownie)
                    logging.getLogger("discord").exception(
                        "Unhandled exception while compiling %s", tex_name
                    )
                    raise

                content = (
                    f"**LaTeX** — {len(tex_names)} plik(ów). Model: `{OPENAI_MODEL}`\n\n"
                    + "\n\n".join(parts)
                )
                await ctx.reply(
                    content if len(content) < 1900 else content[:1900] + "…",
                    files=files_to_send,
                    mention_author=False,
                )

    # -------- /latexclean (PDF only; log -> debug logger) --------
    @commands.hybrid_command(
        nsfw=False,
        name="latexclean",
        description="Kompiluje .tex i odsyła TYLKO PDF (log w debug).",
        guild=None if GUILD_ID is None else discord.Object(id=GUILD_ID),
    )
    @commands.has_any_role(*ALLOWED_ROLES)
    async def latexclean(self, ctx: commands.Context):
        ch = ctx.message.channel
        async with ch.typing():
            all_bytes = await _collect_attachments_bytes(
                ctx.message.attachments, LATEX_MAX_ATTACH_MB
            )
            if not all_bytes:
                return await ctx.reply(
                    f"Dołącz .tex (≤{LATEX_MAX_ATTACH_MB} MiB).", mention_author=False
                )

            # PODZIEL NA GŁÓWNE .TEX i WSPIERAJĄCE
            tex_names = [
                n
                for n in all_bytes.keys()
                if n.lower().endswith(".tex") and is_safe_tex_name(n)
            ]
            if not tex_names:
                return await ctx.reply(
                    f"Dołącz .tex (≤{LATEX_MAX_ATTACH_MB} MiB).", mention_author=False
                )

            support_names = [
                n
                for n in all_bytes.keys()
                if n not in tex_names and is_safe_asset_name(n)
            ]

            files_to_send = []
            parts = []

            for tex_name in tex_names:
                try:
                    main_bytes = all_bytes[tex_name]
                    support = [
                        (n, all_bytes[n])
                        for n in support_names
                        if n != tex_name and is_safe_asset_name(n)
                    ]
                    # (ważne) do supportów dopuszczamy również **inne .tex** – dla \input/\include
                    other_tex = [
                        (n, all_bytes[n])
                        for n in tex_names
                        if n != tex_name and is_safe_tex_name(n)
                    ]
                    support.extend(other_tex)

                    res = await compile_single_tex_bytes(
                        main_bytes,
                        tex_name,
                        LATEX_TEX_ENGINE,
                        LATEX_MAX_COMPILE_SECONDS,
                        logger="discord",
                        support_files=support,
                    )

                    # Twoja dotychczasowa obsługa success/fail – przykładowo:
                    if res["ok"] and res["pdf_bytes"]:
                        b = io.BytesIO(res["pdf_bytes"])
                        b.seek(0)
                        b.name = res["pdf_name"]
                        files_to_send.append(discord.File(b, filename=res["pdf_name"]))
                        parts.append(f"✅ `{tex_name}` → `{res['pdf_name']}`")
                    else:
                        parts.append(
                            f"❌ `{tex_name}` — błąd kompilacji.)"
                        )
                        # Tutaj wywołujesz swoją ścieżkę do OpenAI diagnozy, którą ogarnęliśmy wcześniej.
                except Exception:
                    # let it crash — pełny stack w logu i re-raise (Discord.py to pokaże sensownie)
                    logging.getLogger("discord").exception(
                        "Unhandled exception while compiling %s", tex_name
                    )
                    raise
            await ctx.reply(
                "**LaTeX clean** — " + ", ".join(parts),
                files=files_to_send,
                mention_author=False,
            )

    # -------- /texbatch (ZIP -> ZIP) --------
    @app_commands.command(
        name="texbatch",
        description="Wyślij ZIP z .tex → odeślę ZIP z PDF-ami (tylko udane).",
    )
    @app_commands.describe(archive=f"Archiwum .zip (max ~{LATEX_MAX_ZIP_MB} MiB)")
    @app_commands.checks.has_any_role(*ALLOWED_ROLES)
    async def texbatch(
        self, interaction: discord.Interaction, archive: discord.Attachment
    ):
        await interaction.response.defer()
        if not archive or not archive.filename.lower().endswith(".zip"):
            return await interaction.followup.send(
                "Dołącz ZIP (.zip) z plikami .tex.", ephemeral=True
            )
        if archive.size > LATEX_MAX_ZIP_MB * 1024 * 1024:
            return await interaction.followup.send(
                f"ZIP > {LATEX_MAX_ZIP_MB} MiB — za duży.", ephemeral=True
            )

        data = await archive.read()
        out_zip_bytes, failed = await compile_zip_to_zip(
            data, LATEX_TEX_ENGINE, LATEX_MAX_COMPILE_SECONDS, logger=self.logger_name
        )
        if not out_zip_bytes and failed:
            return await interaction.followup.send(
                "Brak PDF-ów. " + "; ".join(failed[:10]), ephemeral=True
            )

        bio = io.BytesIO(out_zip_bytes)
        bio.seek(0)
        bio.name = "compiled_pdfs.zip"
        if failed:
            self.logger.debug("BATCH failed: %s", ", ".join(failed))
        await interaction.followup.send(
            "✅ Gotowe (ZIP w załączniku)."
            + (f" ❌ Błędy: {len(failed)}" if failed else ""),
            files=[discord.File(bio, filename=bio.name)],
        )


async def setup(bot):
    logger = logging.getLogger("discord")
    await bot.add_cog(LatexModule(bot, logger_name="discord"))
    logger.info("Loading kinky latex module done")
