from transformers import AutoTokenizer, AutoModelForCausalLM, AutoTokenizer
import soundfile as sf
import torch

from datetime import datetime
import random
import time
import numpy as np
import os
import argparse
import gradio as gr


def save_to_txt(text_to_save):
    with open('prompt.txt', 'w', encoding='utf-8') as f:
        f.write(text_to_save)
        
def read_txt():
    with open('prompt.txt') as f:
        lines = f.readlines()
    return lines 



##### Chat z LLAMA ####
##### Chat z LLAMA ####
##### Chat z LLAMA ####


def _load_model_tokenizer():
    model_id = 'tangger/Qwen-7B-Chat'
    tokenizer = AutoTokenizer.from_pretrained(model_id,  trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto",trust_remote_code=True, fp16=True).eval()  
    return model, tokenizer
model_llm, tokenizer_llm = _load_model_tokenizer()


def postprocess(self, y):
    if y is None:
        return []
    for i, (message, response) in enumerate(y):
        y[i] = (
            None if message is None else mdtex2html.convert(message),
            None if response is None else mdtex2html.convert(response),
        )
    return y


def _parse_text(text):
    lines = text.split("\n")
    lines = [line for line in lines if line != ""]
    count = 0
    for i, line in enumerate(lines):
        if "```" in line:
            count += 1
            items = line.split("`")
            if count % 2 == 1:
                lines[i] = f'<pre><code class="language-{items[-1]}">'
            else:
                lines[i] = f"<br></code></pre>"
        else:
            if i > 0:
                if count % 2 == 1:
                    line = line.replace("`", r"\`")
                    line = line.replace("<", "&lt;")
                    line = line.replace(">", "&gt;")
                    line = line.replace(" ", "&nbsp;")
                    line = line.replace("*", "&ast;")
                    line = line.replace("_", "&lowbar;")
                    line = line.replace("-", "&#45;")
                    line = line.replace(".", "&#46;")
                    line = line.replace("!", "&#33;")
                    line = line.replace("(", "&#40;")
                    line = line.replace(")", "&#41;")
                    line = line.replace("$", "&#36;")
                lines[i] = "<br>" + line
    text = "".join(lines)
    return text

def predict(_query, _chatbot, _task_history):
    print(f"User: {_parse_text(_query)}")
    _chatbot.append((_parse_text(_query), ""))
    full_response = ""

    for response in model_llm.chat_stream(tokenizer_llm, _query, history=_task_history,system = "Jesteś assystentem AI. Odpowiadaj zawsze w języku poslkim"  ):
        _chatbot[-1] = (_parse_text(_query), _parse_text(response))

        yield _chatbot
        full_response = _parse_text(response)

    print(f"History: {_task_history}")
    _task_history.append((_query, full_response))
    print(f"Qwen-7B-Chat: {_parse_text(full_response)}")

def read_text(text):
    print("___Tekst do przeczytania!")
    inputs = tokenizer_tss(text, return_tensors="pt").to("cuda")
    with torch.no_grad():
        output = tts_model(**inputs).waveform.squeeze().cpu().numpy()
    sf.write('temp_file.wav', output, tts_model.config.sampling_rate)
    return 'temp_file.wav'


def update_audio(text):
    return 'temp_file.wav'

def translate(audio):
    print("__Wysyłam nagranie do whisper!")
    transcription = whisper_model.transcribe(audio, language="pl")
    return transcription["text"]


def predict(audio, _chatbot, _task_history):
    # Użyj funkcji translate, aby przekształcić audio w tekst
    _query = translate(audio)

    print(f"____User: {_parse_text(_query)}")
    _chatbot.append((_parse_text(_query), ""))
    full_response = "" 

    for response in model_llm.chat_stream(tokenizer_llm,
                                      _query,
                                      history= _task_history,
                                      system = "Jesteś assystentem AI. Odpowiadaj zawsze w języku polskim. Odpowiadaj krótko."):
        _chatbot[-1] = (_parse_text(_query), _parse_text(response))
        yield _chatbot
        full_response = _parse_text(response)

    print(f"____History: {_task_history}")
    _task_history.append((_query, full_response))
    print(f"__Qwen-7B-Chat: {_parse_text(full_response)}")
    print("____full_response",full_response)
    audio_file = read_text(_parse_text(full_response))  # Generowanie audio
    return full_response
    # return 'temp_file.wav'  # Zwrócenie ścieżki do pliku audio

def regenerate(_chatbot, _task_history):
    if not _task_history:
        yield _chatbot
        return
    item = _task_history.pop(-1)
    _chatbot.pop(-1)
    yield from predict(item[0], _chatbot, _task_history)

with gr.Blocks() as chat_demo:
    chatbot = gr.Chatbot(label='Llama Voice Chatbot', elem_classes="control-height")
    query = gr.Textbox(lines=2, label='Input')
    task_history = gr.State([])
    audio_output = gr.Audio('temp_file.wav', label="Generated Audio (wav)", type='filepath', autoplay=False)
    
    with gr.Row():
        submit_btn = gr.Button("🚀 Wyślij tekst")

    with gr.Row():
        audio_upload = gr.Audio(source="microphone", type="filepath", show_label=False)
        submit_audio_btn = gr.Button("🎙️ Wyślij audio")

    submit_btn.click(predict, [query, chatbot, task_history], [chatbot], show_progress=True)
    submit_audio_btn.click(predict, [audio_upload, chatbot, task_history], [chatbot], show_progress=True).then(update_audio, chatbot, audio_output)





##### Audio Gen ####
##### Audio Gen ####
##### Audio Gen ####


##### Run Alll #######
##### Run Alll #######
##### Run Alll #######


demo_all = gr.TabbedInterface([chat_demo], ["Chat with LLama"])

demo_all.queue()
demo_all.launch()