import discord
import logging
import os
from huggingface_hub import InferenceClient
import asyncio
import subprocess
from datasets import load_dataset
import pandas as pd
from fuzzywuzzy import process

# 현재 작업 디렉토리 출력
print("Current Working Directory:", os.getcwd())

# 데이터셋 파일 이름
data_files = ['train_0.csv', 'train_1.csv', 'train_2.csv', 'train_3.csv', 'train_4.csv', 'train_5.csv']

# 현재 작업 디렉토리에 모든 파일이 있는지 확인
missing_files = [file for file in data_files if not os.path.exists(file)]
if missing_files:
    print(f"Missing files: {missing_files}")
    # 필요한 경우 작업 디렉토리 변경
    os.chdir('/home/user/app')
    print("Changed directory to:", os.getcwd())
else:
    print("All files are present in the current directory.")

# 데이터셋 로드 및 최적화
def load_optimized_dataset(data_files):
    data_frames = [pd.read_csv(file) for file in data_files]
    full_data = pd.concat(data_frames, ignore_index=True)
    # 사건명을 키로 하고 사건번호와 전문을 저장하는 딕셔너리 생성
    name_to_number = full_data.groupby('사건명')['사건번호'].apply(list).to_dict()
    number_to_fulltext = full_data.set_index('사건번호')['전문'].to_dict()
    return name_to_number, number_to_fulltext

name_to_number, number_to_fulltext = load_optimized_dataset(data_files)
print("Dataset loaded successfully.")

# 사건명 리스트 생성
all_case_names = list(name_to_number.keys())

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

# 특정 채널 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:
            logging.debug("Currently processing another message, skipping this one.")
            return

        self.is_processing = True
        try:
            response = await generate_response(message)
            if response and response.strip():
                await message.channel.send(response)
            else:
                await message.channel.send("죄송합니다, 제공할 수 있는 정보가 없습니다.")
        finally:
            self.is_processing = False
            logging.debug("Message processing completed, ready for the next one.")

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

async def generate_response(message):
    global conversation_history
    user_input = message.content.strip()
    user_mention = message.author.mention

    # 유사한 사건명 찾기
    matched_case_names = process.extractBests(user_input, all_case_names, limit=3, score_cutoff=80)
    
    if matched_case_names:
        case_numbers = []
        for case_name, score in matched_case_names:
            case_numbers.extend(name_to_number[case_name])
        case_numbers = list(set(case_numbers))  # 중복 제거
        case_numbers_str = "\n".join(case_numbers)
        system_message = f"{user_mention}, '{user_input}'와 유사한 사건명의 사건번호는 다음과 같습니다:\n{case_numbers_str}"
    elif user_input in number_to_fulltext:
        full_text = number_to_fulltext[user_input]
        system_message = f"{user_mention}, 사건번호 '{user_input}'의 전문은 다음과 같습니다:\n\n{full_text}"
    else:
        system_message = f"{user_mention}, 관련 법률 정보를 찾을 수 없습니다."

    # 메시지 길이 제한 처리
    max_length = 2000
    if len(system_message) > max_length:
        response_parts = []
        for i in range(0, len(system_message), max_length):
            part_response = system_message[i:i + max_length]
            await message.channel.send(part_response)
            response_parts.append(part_response)
        return response_parts[0] if response_parts else "죄송합니다, 제공할 수 있는 정보가 없습니다."
    return system_message

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