import discord
import logging
import os
from huggingface_hub import InferenceClient
from transformers import pipeline
import asyncio
import subprocess
import re

# 로깅 설정
logging.basicConfig(level=logging.DEBUG, format='%(asctime)s:%(levelname)s:%(name)s: %(message)s', handlers=[logging.StreamHandler()])

# 인텐트 설정
intents = discord.Intents.default()
intents.message_content = True
intents.messages = True
intents.guilds = True
intents.guild_messages = True

# 추론 API 클라이언트 설정
hf_client = InferenceClient("CohereForAI/c4ai-command-r-plus", token=os.getenv("HF_TOKEN"))

# 수학 전문 LLM 파이프라인 설정
math_pipe = pipeline("text-generation", model="AI-MO/NuminaMath-7B-TIR")

# 특정 채널 ID
SPECIFIC_CHANNEL_ID = int(os.getenv("DISCORD_CHANNEL_ID"))

# 대화 히스토리를 저장할 전역 변수
conversation_history = []

class MyClient(discord.Client):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.is_processing = False

    async def on_ready(self):
        logging.info(f'{self.user}로 로그인되었습니다!')
        subprocess.Popen(["python", "web.py"])
        logging.info("Web.py server has been started.")

    async def on_message(self, message):
        if message.author == self.user:
            return
        if not self.is_message_in_specific_channel(message):
            return
        if self.is_processing:
            return

        self.is_processing = True
        try:
            if self.is_math_question(message.content):
                response = await self.handle_math_question(message.content)
            else:
                response = await self.generate_response(message)
            await self.send_long_message(message.channel, response)
        finally:
            self.is_processing = False

    def is_message_in_specific_channel(self, message):
        return message.channel.id == SPECIFIC_CHANNEL_ID or (
            isinstance(message.channel, discord.Thread) and message.channel.parent_id == SPECIFIC_CHANNEL_ID
        )

    def is_math_question(self, content):
        return bool(re.search(r'\b(solve|equation|calculate|math)\b', content, re.IGNORECASE))

    async def handle_math_question(self, question):
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(None, lambda: math_pipe([{"role": "user", "content": question}]))
        return response[0]['generated_text']

    async def generate_response(self, message):
        global conversation_history
        user_input = message.content
        user_mention = message.author.mention
        conversation_history.append({"role": "user", "content": user_input})
        messages = [{"role": "system", "content": "한글로 답변하십시오."}] + conversation_history
        response = await asyncio.get_event_loop().run_in_executor(None, lambda: hf_client.chat_completion(
            messages, max_tokens=1000, stream=True, temperature=0.7, top_p=0.85))
        full_response = ''.join([part.choices[0].delta.content for part in response if part.choices and part.choices[0].delta and part.choices[0].delta.content])
        conversation_history.append({"role": "assistant", "content": full_response})
        return f"{user_mention}, {full_response}"

    async def send_long_message(self, channel, message):
        if len(message) <= 2000:
            await channel.send(message)
        else:
            parts = [message[i:i+2000] for i in range(0, len(message), 2000)]
            for part in parts:
                await channel.send(part)

if __name__ == "__main__":
    discord_client = MyClient(intents=intents)
    discord_client.run(os.getenv('DISCORD_TOKEN'))