import discord
from discord.ext import commands
import logging
import asyncio
import json
import os
import aiohttp
from datetime import datetime, date

logger = logging.getLogger(__name__)


class AIClient(commands.Cog):
    def __init__(self, bot):
        self.bot = bot
        self.config_path = os.path.join(os.path.dirname(__file__), "config.json")
        self.usage_path = os.path.join(os.path.dirname(__file__), "usage.json")
        self.config = self._load_config()
        self.usage = self._load_usage()

        # user_id (int) -> list of {"role": ..., "content": ...}
        self.conversations: dict[int, list[dict]] = {}
        # bot message_id -> user_id that owns that conversation thread
        self.message_map: dict[int, int] = {}
        # user_id -> datetime of last request (for cooldown)
        self.last_used: dict[int, datetime] = {}

    # confeeaggeeee

    def _load_config(self) -> dict:
        with open(self.config_path, "r", encoding="utf-8") as f:
            return json.load(f)

    def _save_config(self) -> None:
        with open(self.config_path, "w", encoding="utf-8") as f:
            json.dump(self.config, f, indent=4, ensure_ascii=False)

    def _load_usage(self) -> dict:
        if os.path.exists(self.usage_path):
            with open(self.usage_path, "r", encoding="utf-8") as f:
                return json.load(f)
        return {}

    def _save_usage(self) -> None:
        with open(self.usage_path, "w", encoding="utf-8") as f:
            json.dump(self.usage, f, indent=4)

    # helpers

    def _get_user_usage(self, user_id: int) -> dict:
        """Return today's usage record for user, resetting if the date changed."""
        uid = str(user_id)
        today = str(date.today())
        if uid not in self.usage or self.usage[uid].get("date") != today:
            self.usage[uid] = {"date": today, "requests": 0, "tokens": 0}
        return self.usage[uid]

    def _check_channel(self, channel_id: int) -> bool:
        allowed = self.config.get("ALLOWED_CHANNEL_IDS", [])
        return (not allowed) or (channel_id in allowed)

    def _cooldown_remaining(self, user_id: int) -> float:
        cooldown = self.config.get("COOLDOWN_SECONDS", 10)
        last = self.last_used.get(user_id)
        if last is None:
            return 0.0
        elapsed = (datetime.now() - last).total_seconds()
        remaining = cooldown - elapsed
        return max(remaining, 0.0)

    def _trim_history(self, user_id: int) -> None:
        max_pairs = self.config.get("MAX_HISTORY_LENGTH", 10)
        history = self.conversations.get(user_id, [])
        cap = max_pairs * 2
        if len(history) > cap:
            self.conversations[user_id] = history[-cap:]

    def _register_bot_message(self, message_id: int, user_id: int) -> None:
        self.message_map[message_id] = user_id
        if len(self.message_map) > 1000:
            for old_id in list(self.message_map.keys())[:200]:
                del self.message_map[old_id]


    async def _call_openrouter(self, messages: list[dict]) -> tuple[str, int]:
        # call api
        api_key = self.config.get("OPENROUTER_API_KEY", "")
        model = self.config.get("MODEL", "openai/gpt-4o-mini")
        max_tokens = self.config.get("MAX_RESPONSE_TOKENS", 1000)

        headers = {
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "HTTP-Referer": "https://discord.bot",
            "X-Title": "Discord Bot",
        }
        payload = {
            "model": model,
            "messages": messages,
            "max_tokens": max_tokens,
        }

        async with aiohttp.ClientSession() as session:
            async with session.post(
                "https://openrouter.ai/api/v1/chat/completions",
                headers=headers,
                json=payload,
                timeout=aiohttp.ClientTimeout(total=60),
            ) as resp:
                if resp.status != 200:
                    body = await resp.text()
                    raise RuntimeError(f"HTTP {resp.status}: {body[:300]}")
                data = await resp.json()

        text = data["choices"][0]["message"]["content"]
        tokens = data.get("usage", {}).get("total_tokens", 0)
        return text, tokens

    async def _process_request(
        self,
        message: discord.Message,
        user_input: str,
    ) -> None:
        # logic for !ai
        author = message.author
        channel = message.channel

        if not self._check_channel(channel.id):
            allowed_ids = self.config.get("ALLOWED_CHANNEL_IDS", [])
            mention_list = ", ".join(f"<#{cid}>" for cid in allowed_ids)
            await channel.send(
                f"Ai is only available in: {mention_list}",
                delete_after=10,
            )
            return

        remaining = self._cooldown_remaining(author.id)
        if remaining > 0:
            await channel.send(
                f"Hit cooldown. Wait for **{remaining:.1f}s**.",
                delete_after=10,
            )
            return

        usage = self._get_user_usage(author.id)
        max_req = self.config.get("MAX_REQUESTS_PER_DAY", 50)
        max_tok = self.config.get("MAX_TOKENS_PER_USER_PER_DAY", 10000)

        if usage["requests"] >= max_req:
            await channel.send(
                f"{author.mention}, daily limit hit "
                f"({max_req}). Try tommorow.",
                delete_after=15,
            )
            return

        if usage["tokens"] >= max_tok:
            await channel.send(
                f"{author.mention}, daily token limit hit "
                f"({max_tok}). Try tommorow.",
                delete_after=15,
            )
            return

        if author.id not in self.conversations:
            self.conversations[author.id] = []

        self.conversations[author.id].append({"role": "user", "content": user_input})
        self._trim_history(author.id)

        system_prompt = self.config.get(
            "SYSTEM_PROMPT", "You are a helpful assistant in a Discord server."
        )
        api_messages = [{"role": "system", "content": system_prompt}] + self.conversations[author.id]

        self.last_used[author.id] = datetime.now()

        async with channel.typing():
            try:
                response_text, tokens_used = await self._call_openrouter(api_messages)
            except Exception as exc:
                logger.error("OpenRouter error: %s", exc)
                self.conversations[author.id].pop()
                await channel.send(f"err: `{exc}`")
                return

        self.conversations[author.id].append(
            {"role": "assistant", "content": response_text}
        )

        usage["requests"] += 1
        usage["tokens"] += tokens_used
        self._save_usage()

        chunks = [response_text[i: i + 1990] for i in range(0, len(response_text), 1990)]
        last_sent: discord.Message | None = None
        for idx, chunk in enumerate(chunks):
            if idx == 0:
                last_sent = await message.reply(chunk)
            else:
                last_sent = await channel.send(chunk)

        if last_sent:
            self._register_bot_message(last_sent.id, author.id)

    # cmds

    @commands.command(name="ai")
    async def ai(self, ctx: commands.Context, *, prompt: str):
        # i really wonder what this does
        await self._process_request(ctx.message, prompt)

    @commands.command(name="ai_reset")
    async def ai_reset(self, ctx: commands.Context):
        # clear your personal conversation history with the ai
        self.conversations.pop(ctx.author.id, None)
        await ctx.reply("Reset personal conversation history.")

    #!! admin cmds

    @commands.command(name="ai_setchannel")
    @commands.has_permissions(administrator=True)
    async def ai_setchannel(
        self, ctx: commands.Context, channel: discord.TextChannel = None
    ):
        # whitelist a channel for ai commands, or clear the whitelist if no channel is given
        if channel is None:
            self.config["ALLOWED_CHANNEL_IDS"] = []
            self._save_config()
            await ctx.reply("Whitelist is not active")
            return

        ids: list[int] = self.config.setdefault("ALLOWED_CHANNEL_IDS", [])
        if channel.id in ids:
            await ctx.reply(f"{channel.mention} is already in whitelist")
        else:
            ids.append(channel.id)
            self._save_config()
            await ctx.reply(f"Added {channel.mention} into whitelist.")

    @commands.command(name="ai_removechannel")
    @commands.has_permissions(administrator=True)
    async def ai_removechannel(self, ctx: commands.Context, channel: discord.TextChannel):
        # remove channel from whitelist
        ids: list[int] = self.config.get("ALLOWED_CHANNEL_IDS", [])
        if channel.id not in ids:
            await ctx.reply(f"{channel.mention} is not in whitelist")
        else:
            ids.remove(channel.id)
            self._save_config()
            await ctx.reply(f"Removed {channel.mention} from whitelist")

    @commands.command(name="ai_clearlimit")
    @commands.has_permissions(administrator=True)
    async def ai_clearlimit(self, ctx: commands.Context, member: discord.Member):
        # reset limit for a specific member
        uid = str(member.id)
        if uid in self.usage:
            del self.usage[uid]
            self._save_usage()
        await ctx.reply(f"Reset limits for {member.mention}")

    @commands.command(name="ai_globalstats")
    @commands.has_permissions(administrator=True)
    async def ai_globalstats(self, ctx: commands.Context):
        # todays usage
        today = str(date.today())
        lines: list[str] = []
        total_req = 0
        total_tok = 0

        for uid, data in self.usage.items():
            if data.get("date") != today:
                continue
            member = ctx.guild.get_member(int(uid))
            name = member.display_name if member else f"User {uid}"
            req = data["requests"]
            tok = data["tokens"]
            total_req += req
            total_tok += tok
            lines.append(f"**{name}** - {req} requests, {tok} tokens")

        if not lines:
            await ctx.reply("No water wasted yet.")
            return

        lines.sort()
        embed = discord.Embed(
            title=f"How much water wasted for the date of {today}",
            description="\n".join(lines),
            color=discord.Color.orange(),
        )
        await ctx.reply(embed=embed)

    # reply

    @commands.Cog.listener()
    async def on_message(self, message: discord.Message):
        # ignore bots
        if message.author.bot:
            return

        # only handle messages that are a reply to another message
        if not message.reference:
            return

        ctx = await self.bot.get_context(message)
        if ctx.valid:
            return

        
        ref = message.reference.resolved
        if not isinstance(ref, discord.Message):
            # Not cached — fetch it
            try:
                ref = await message.channel.fetch_message(message.reference.message_id)
            except (discord.NotFound, discord.HTTPException):
                return

        # must be a reply to the bot
        if ref.author.id != self.bot.user.id:
            return

        # must be a message we registered as part of a conversation
        if ref.id not in self.message_map:
            return

        await self._process_request(message, message.content)

    # utility

    @staticmethod
    def _progress_bar(value: int, maximum: int, length: int = 10) -> str:
        if maximum == 0:
            return "█" * length
        filled = round(length * min(value, maximum) / maximum)
        return "█" * filled + "░" * (length - filled)


async def setup(bot):
    await bot.add_cog(AIClient(bot))