import os
import sys
import ast
import csv

import numpy as np
from IPython.display import display
from dotenv import load_dotenv
from scipy import spatial
import pandas as pd
import tiktoken

from .openai_connect import get_openai_client
from datetime import datetime
 
load_dotenv() 

EMBEDDING_MODEL = "text-embedding-3-small"
EMBEDDING_ENCODING = "cl100k_base"
MAX_TOKENS = 8000
GPT_MODEL = os.getenv("GPT_MODEL")
 

def strings_ranked_by_relatedness(
        query:str,
        df:pd.DataFrame,
        relatedness_fn = lambda x,y: 1 - spatial.distance.cosine(x,y),
        top_n:int = 100,
        show_relatednesses:bool = False,
        additional_columns=[]
):
    client = get_openai_client()
    query_embedding_response = client.embeddings.create(
        input = [query],
        model = EMBEDDING_MODEL
    )
    query_embedding = query_embedding_response.data[0].embedding

    strings_and_relatednesses = [
      (row["combined"], relatedness_fn(query_embedding, row["ada_embedding"])) + tuple(row[col] for col in additional_columns)
      for i, row in df.iterrows()
    ]

    filtered_strings_and_relatednesses = [
        item for item in strings_and_relatednesses if item[1] >= 0.6
    ]

    filtered_strings_and_relatednesses.sort(key=lambda x:x[1],reverse=True)

    if not filtered_strings_and_relatednesses:
        return tuple([] for _ in range(2 + len(additional_columns)))
    
    results = zip(*filtered_strings_and_relatednesses)

    results = list(results)
    results[-1] = tuple(f"{q} 相關係數: {round(rel, 3)}" for q, rel in zip(results[-1], results[1]))
    print(f"相關係數: {results[1]}")

    
    return tuple(result[:top_n] for result in results)


def num_tokens(text:str,model:str = GPT_MODEL):
    encoding = tiktoken.encoding_for_model(model)
    return len(encoding.encode(text))

def query_message(
    query:str,
    df:pd.DataFrame,
    model:str,
    token_budget:int,
    introduction:str = "",
) -> str:
    strings, relatednesses = strings_ranked_by_relatedness(query, df, top_n=5,show_relatednesses=True)
    introduction = introduction
    question = f"{query}"
    message = introduction
    print(f"相關係數: {relatednesses}")

    if len(strings) == 0:
        return f"無法辨識"
    
    for string in strings:
        next_article = f'\n\nWikipedia article section:\n"""\n{string}\n"""'
        if (
            num_tokens(message + next_article + question, model=model)
            > token_budget
        ):
            break
        else:
            message += next_article
    return message + question

def get_embedding(text:str,model:str = EMBEDDING_MODEL):
    client = get_openai_client()
    text = text.replace("\n"," ")
    return client.embeddings.create(
        input = [text],
        model = model
    ).data[0].embedding

def load_or_create_maintain_embedding():
    if os.path.exists("data/維護訊息_embedding.csv"):
        df = pd.read_csv("data/維護訊息_embedding.csv")
        df["ada_embedding"] = df.ada_embedding.apply(ast.literal_eval)
        return df
    else:
        df = pd.read_csv('data/維護訊息22.csv')
        df = df[['類型','訊息內容','遊戲平台']]
        df["combined"] = (
            "問題類型為: " + df["類型"] + "\n\n\n\n" +
            "內容訊息為: " + df["訊息內容"] + "\n\n\n\n" +
            "遊戲平台為: " + df["遊戲平台"]
        )

        top_n = 100
        encoding = tiktoken.encoding_for_model(EMBEDDING_MODEL)
        df["n_tokens"] = df.combined.apply(lambda x: len(encoding.encode(x)))
        df = df[df.n_tokens <= MAX_TOKENS].tail(top_n)
        df["ada_embedding"] = df.combined.apply(lambda x: get_embedding(x,EMBEDDING_MODEL))
        df.to_csv("data/維護訊息_embedding.csv",index=False)
        print("!!")
        return df


def ask(
    query:str,
    model:str = GPT_MODEL,
    token_budget: int = 4096 - 500,
    print_message:bool = False,
) -> str:
    df = load_or_create_maintain_embedding()
    client = get_openai_client()
    now = datetime.now()
    fmt_time_date = now.strftime("%Y-%m-%d")
    fmt_time_time = now.strftime("%H:%M:%S")
    print(query)
    message = query_message(query,df,model=model,token_budget=token_budget,introduction=f"""
            使用這份csv的“訊息內容”，
            先判斷出“訊息內容”的類型是什麼，
            訊息內沒有提及”維護“，以及符合”遊戲平台“的字眼的話，那他通常不是維護訊息，
            請你自行判斷{query}裡面的維護時間，不要直接回傳csv檔案裡面的訊息，那只是給你參考用的
            而不是csv裡面的時間，
            如果類型為“維護訊息”請遵循以下規定：
                找出使用者回傳最相關的那列，
                並且將“訊息內容”整理出一個資訊：遊戲平台
                請注意，
                我要你自己從用戶傳的訊息判斷出“開始維護”和“結束維護”是什麼時候，
                如果內容有多個時間的話，關鍵字要找“遊戲”之類的，
                我提供你一些建議：
                    1. 03:00 a.m. - 12:00 p.m. 的維護時間代表 03:00:00 - 12:00:00
                    2. 05/02 (一) 1:00 的維護時間代表 2025-05-02 01:00:00
                ----    
                如果內容沒有開始日期的資訊，
                請自動將時間設為{fmt_time_date}，格式為YYYY-MM-DD
                如果內容沒有結束日期的資訊，且訊息內容沒有描述結束時間的資訊，
                請自動將時間設為{fmt_time_date}+1天，格式為YYYY-MM-DD，
                請注意，如果有明確的結束時間，通常結束日期會跟開始日期相同，
                例如：
                ※ 5/2 星期一 02:00-20:30，那麼結束日期就是2025-05-02
                ※ 5/2(一) 05:00AM-14:30PM，那麼結束日期就是2025-05-02
                ----
                如果內容沒有開始時間的資訊，
                請自動將時間設為{fmt_time_time}，格式為HH:mm:ss
                如果內容沒有結束時間的資訊，
                請自動將時間設為{fmt_time_time}+1天，格式為HH:mm:ss
                                
                回傳“類型”,“遊戲平台”,“開始日期”,“開始時間”,“結束日期“,”結束時間“,“用戶傳過來的訊息”
                例如:
                    訊息種類：[類型]\n
                    遊戲平台：ATG\n
                    開始日期：2025-10-01\n
                    開始時間：00:00:00\n
                    結束日期：2025-10-02\n
                    結束時間：00:00:00\n
                                
                    ----------------------------
                    原始訊息：{query}
                            
            如果類型為你沒有看過的資訊請遵循以下規定：
                回傳
                例如:
                    訊息種類：未知\n
                                
                    ----------------------------
                    原始訊息：[“用戶回傳的訊息”]           
                """)
    if print_message:
      print(message)

    if message == "無法辨識":
        message = "無法回答"
    messages = [
      {"role": "system", "content": """
            請根據使用者送出的訊息回答最相關的答案，
            如果訊息相關度不高或是具有冒犯性的話，
            請回覆“無法回答”。
       """},
      {"role": "user", "content": message},
    ]
    response = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=0
    )
    response_message = response.choices[0].message.content
    print(response_message)
    return response_message


 
