import discord
import logging
import os
import json
from huggingface_hub import InferenceClient
import asyncio
import subprocess
from sentence_transformers import SentenceTransformer, util
import torch

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

# JSON 데이터셋 로드
with open("jangtest.json", "r", encoding="utf-8") as f:
    dataset = json.load(f)

# 문장 임베딩 모델 로드
model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')

# 데이터셋의 임베딩을 미리 계산
dataset_texts = [json.dumps(item, ensure_ascii=False) for item in dataset]
dataset_embeddings = model.encode(dataset_texts, convert_to_tensor=True)

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
    
    # 유사한 데이터 찾기
    most_similar_data = find_most_similar_data(user_input)
    
    system_message = f"""
    당신은 'kAI'라는 이름의 한국 보험 상품에 대한 AI 조언자입니다. 
    반드시 제공된 데이터셋의 정보만을 사용하여 답변해야 합니다. 
    제공된 데이터에 없는 정보에 대해서는 "죄송합니다. 해당 정보는 제가 가진 데이터에 없습니다."라고 답변하십시오.
    모든 답변은 한글로 하고, markdown 형식으로 출력하세요.
    """
    
    conversation_history.append({"role": "user", "content": user_input})
    messages = [{"role": "system", "content": system_message}] + conversation_history
    
    if most_similar_data:
        messages.append({"role": "system", "content": f"관련 정보: {most_similar_data}"})
    else:
        return f"{user_mention}, 죄송합니다. 귀하의 질문과 관련된 정보를 찾을 수 없습니다."
    
    logging.debug(f'Messages to be sent to the model: {messages}')
    
    loop = asyncio.get_event_loop()
    response = await 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 = []
    for part in response:
        logging.debug(f'Part received from stream: {part}')
        if part.choices and part.choices[0].delta and part.choices[0].delta.content:
            full_response.append(part.choices[0].delta.content)
    
    full_response_text = ''.join(full_response)
    logging.debug(f'Full model response: {full_response_text}')
    
    conversation_history.append({"role": "assistant", "content": full_response_text})
    return f"{user_mention}, {full_response_text}"

def find_most_similar_data(query):
    query_embedding = model.encode(query, convert_to_tensor=True)
    
    # 코사인 유사도 계산
    cos_scores = util.pytorch_cos_sim(query_embedding, dataset_embeddings)[0]
    top_result = torch.topk(cos_scores, k=1)
    
    if top_result.values[0] > 0.5:  # 임계값 설정
        return json.dumps(dataset[top_result.indices[0]], ensure_ascii=False, indent=2)
    else:
        return None

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