from openai import OpenAI
from api.openai_scripts_new.config import SYSTEM_PROMPT, OPENAI_API_KEY, SUPABASE_KEY, SUPABASE_URL
from supabase import create_client, Client
from api.openai_scripts_new.text_processing import fuzzy_correct_chinese
from transformers import pipeline
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.callbacks import get_openai_callback
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
import torchaudio
import torch

client = OpenAI(api_key=OPENAI_API_KEY)
supabase: Client = create_client(SUPABASE_URL, SUPABASE_KEY)

pipe = pipeline(model="linshoufan/linshoufanfork-whisper-small-nan-tw-pinyin")

def transcribe(audio_data):
    try:
        # table_name = "word_database"
        # response = supabase.table(table_name).select("term").execute()
        # custom_vocab = []
        # if response.data:
        #     for item in response.data:
        #         custom_vocab.append({item['term']})
        # else:
        #     print(f"No data found or an error occurred: {response.error}")
        #     print("Using default dictionary as Supabase data couldn't be fetched.")
        # 如果音頻是立體聲，轉換為單聲道
        text = pipe(audio_data)["text"]
        print(transcript)
        # 翻譯台羅拼音為繁體中文
        model_name = "gpt-4o"
        llm = ChatOpenAI(model_name=model_name, temperature=0.7, api_key=OPENAI_API_KEY, max_tokens=4096)
        with get_openai_callback() as cb:
            qa_system_prompt = f"""你是一個專門翻譯台羅拼音的助理,可以將台語音精準的轉換成在繁體中文中的意思.
    你是一名資深的大語言模型領域的專家,精通模型架構原理和落地應用實踐,只需要翻譯成繁體中文即可."""

            qa_prompt = ChatPromptTemplate.from_messages(
                [
                    ("system", qa_system_prompt),
                    ("human", "{transcript}"),
                ]
            )

            rag_chain = (
                qa_prompt
                | llm
                | StrOutputParser()
            )

            # session_id = "abc123"  # 這應該是從某個上下文獲取的動態值
            # chat_history = get_session_history(session_id)

            text = rag_chain.invoke(
                {"transcript": transcript}
            )

            # # 更新聊天歷史
            # chat_history.add_user_message(inp)
            # chat_history.add_ai_message(text)
            # # chat_history.add_message({'role':HumanMessage(content=input), 'message':AIMessage(content=text)})
            # save_session_history(session_id, chat_history)
            print(f"Total Tokens: {cb.total_tokens}")
            print(f"Prompt Tokens: {cb.prompt_tokens}")
            print(f"Completion Tokens: {cb.completion_tokens}")
            print(f"Total Cost (USD): ${cb.total_cost}")

        return text
    except Exception as e:
        print(f"轉錄時發生錯誤：{str(e)}")
        return None

def post_process_transcript(transcript, temperature=0):
    corrected_transcript = fuzzy_correct_chinese(transcript)
    
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": f"請校對並修正以下轉錄文本，但不要改變其原意或回答問題：\n\n{corrected_transcript}"}
    ]

    response = client.chat.completions.create(
        model="gpt-4",
        temperature=temperature,
        messages=messages
    )

    return response.choices[0].message.content

def process_audio(audio_data):
    raw_transcript = transcribe(audio_data)
    print(raw_transcript)
    if raw_transcript is None:
        return None, None
    corrected_transcript = post_process_transcript(raw_transcript)
    return raw_transcript, corrected_transcript