import discord
import logging
import os
from huggingface_hub import InferenceClient
import asyncio
import subprocess
from datasets import load_dataset

# 로깅 설정
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

# 데이터셋 로드
data_files = ['train_0.csv', 'train_1.csv', 'train_2.csv', 'train_3.csv', 'train_4.csv', 'train_5.csv']
law_dataset = load_dataset('csv', data_files=data_files)

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

# 특정 채널 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:
            response = await generate_response(message)
            await message.channel.send(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
        )

async def generate_response(message):
    global conversation_history
    user_input = message.content
    user_mention = message.author.mention
    system_message = f"{user_mention}, DISCORD에서 사용자들의 질문에 답하는 어시스턴트입니다."
    answer = search_in_dataset(user_input, law_dataset)
    full_response_text = system_message + "\n\n" + answer
    await message.channel.send(full_response_text)
    conversation_history.append({"role": "assistant", "content": full_response_text})

def search_in_dataset(query, dataset):
    response = []
    for record in dataset['train']:
        if record['사건명'] and query in record['사건명']:
            # 사건명으로 검색 시 사건번호만 반환
            response.append(f"사건번호: {record['사건번호']}")
        elif record['사건번호'] == query:
            # 사건번호로 검색 시 전체 내용 반환
            detail = (
                f"판례정보일련번호: {record['판례정보일련번호']}\n"
                f"사건명: {record['사건명']}\n"
                f"사건번호: {record['사건번호']}\n"
                f"선고일자: {record['선고일자']}\n"
                f"선고: {record['선고']}\n"
                f"법원명: {record['법원명']}\n"
                f"사건종류명: {record['사건종류명']}\n"
                f"판결유형: {record['판결유형']}\n"
                f"판시사항: {record['판시사항']}\n"
                f"판결요지: {record['판결요지']}\n"
                f"참조조문: {record['참조조문']}\n"
                f"참조판례: {record['참조판례']}\n"
                f"전문: {record['전문']}\n"
            )
            response.append(detail)
    return "\n".join(response) if response else "관련 법률 정보를 찾을 수 없습니다."


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