import csv
import datetime
import os
import re
import subprocess
import time
import uuid
from io import BytesIO, StringIO
import gradio as gr
import spaces
import torch
import torchaudio
from huggingface_hub import HfApi, hf_hub_download, snapshot_download
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from vinorm import TTSnorm
from content_generation import create_content  # Nhập hàm create_content từ file content_generation.py
from sports_news import get_sports_news_content  # Nhập hàm lấy tin thể thao
from PIL import Image
from pathlib import Path
import requests
import json
import hashlib

# Load environment variables
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    print("Warning: python-dotenv not installed. Using system environment variables only.")

# Download for mecab
os.system("python -m unidic download")

# Cấu hình API và mô hình
HF_TOKEN = os.environ.get("HF_TOKEN")
api = HfApi(token=HF_TOKEN)

# Tải mô hình viXTTS
print("Downloading if not downloaded viXTTS")
checkpoint_dir = "model/"
repo_id = "capleaf/viXTTS"
use_deepspeed = False
os.makedirs(checkpoint_dir, exist_ok=True)
required_files = ["model.pth", "config.json", "vocab.json", "speakers_xtts.pth"]
files_in_dir = os.listdir(checkpoint_dir)
if not all(file in files_in_dir for file in required_files):
    snapshot_download(
        repo_id=repo_id,
        repo_type="model",
        local_dir=checkpoint_dir,
    )
    hf_hub_download(
        repo_id="coqui/XTTS-v2",
        filename="speakers_xtts.pth",
        local_dir=checkpoint_dir,
    )

xtts_config = os.path.join(checkpoint_dir, "config.json")
config = XttsConfig()
config.load_json(xtts_config)
MODEL = Xtts.init_from_config(config)
MODEL.load_checkpoint(
    config, checkpoint_dir=checkpoint_dir, use_deepspeed=use_deepspeed
)
if torch.cuda.is_available():
    MODEL.cuda()

supported_languages = config.languages
if "vi" not in supported_languages:
    supported_languages.append("vi")

# Hàm chuẩn hóa văn bản tiếng Việt
def normalize_vietnamese_text(text):
    text = (
        TTSnorm(text, unknown=False, lower=False, rule=True)
        .replace("..", ".")
        .replace("!.", "!")
        .replace("?.", "?")
        .replace(" .", ".")
        .replace(" ,", ",")
        .replace('"', "")
        .replace("'", "")
        .replace("AI", "Ây Ai")
        .replace("A.I", "Ây Ai")
        .replace("%", "phần trăm")
    )
    return text

# Hàm tính toán độ dài giữ lại cho audio ngắn
def calculate_keep_len(text, lang):
    """Simple hack for short sentences"""
    if lang in ["ja", "zh-cn"]:
        return -1
    word_count = len(text.split())
    num_punct = text.count(".") + text.count("!") + text.count("?") + text.count(",")
    if word_count < 5:
        return 15000 * word_count + 2000 * num_punct
    elif word_count < 10:
        return 13000 * word_count + 2000 * num_punct
    return -1

# Hàm tạo mô tả ảnh từ nội dung audio
def generate_image_description(prompt):
    return f"A visual representation of: {prompt}"

# Hàm gọi API tạo ảnh
def txt2img(prompt, width, height):
    model_id = "770694094415489962"  # Model ID cố định
    vae_id = "sdxl-vae-fp16-fix.safetensors"  # VAE cố định
    lora_items = [
        {"loraModel": "766419665653268679", "weight": 0.7},
        {"loraModel": "777630084346589138", "weight": 0.7},
        {"loraModel": "776587863287492519", "weight": 0.7}
    ]
    txt2img_data = {
        "request_id": hashlib.md5(str(int(time.time())).encode()).hexdigest(),
        "stages": [
            {
                "type": "INPUT_INITIALIZE",
                "inputInitialize": {
                    "seed": -1,
                    "count": 1
                }
            },
            {
                "type": "DIFFUSION",
                "diffusion": {
                    "width": width,
                    "height": height,
                    "prompts": [
                        {
                            "text": prompt
                        }
                    ],
                    "negativePrompts": [
                        {
                            "text": "nsfw"
                        }
                    ],
                    "sdModel": model_id,
                    "sdVae": vae_id,
                    "sampler": "Euler a",
                    "steps": 20,
                    "cfgScale": 3,
                    "clipSkip": 1,
                    "etaNoiseSeedDelta": 31337,
                    "lora": {
                        "items": lora_items
                    }
                }
            }
        ]
    }
    body = json.dumps(txt2img_data)
    headers = {
        'Content-Type': 'application/json',
        'Accept': 'application/json',
        'Authorization': f'Bearer {os.getenv("api_key_token")}'
    }
    response = requests.post(f"https://ap-east-1.tensorart.cloud/v1/jobs", json=txt2img_data, headers=headers)
    if response.status_code != 200:
        return f"Error: {response.status_code} - {response.text}"
    response_data = response.json()
    job_id = response_data['job']['id']
    print(f"Job created. ID: {job_id}")
    start_time = time.time()
    timeout = 300  # Giới hạn thời gian chờ là 300 giây (5 phút)
    while True:
        time.sleep(10)
        elapsed_time = time.time() - start_time
        if elapsed_time > timeout:
            return f"Error: Job timed out after {timeout} seconds."
        response = requests.get(f"https://ap-east-1.tensorart.cloud/v1/jobs/{job_id}", headers=headers)
        if response.status_code != 200:
            return f"Error: {response.status_code} - {response.text}"
        get_job_response_data = response.json()
        job_status = get_job_response_data['job']['status']
        print(f"Job status: {job_status}")
        if job_status == 'SUCCESS':
            if 'successInfo' in get_job_response_data['job']:
                image_url = get_job_response_data['job']['successInfo']['images'][0]['url']
                print(f"Job succeeded. Image URL: {image_url}")
                response_image = requests.get(image_url)
                img = Image.open(BytesIO(response_image.content))
                return img
            else:
                return "Error: Output is missing in the job response."
        elif job_status == 'FAILED':
            return "Error: Job failed. Please try again with different settings."

# Hàm tạo video từ ảnh và audio
def create_video(image_path, audio_path, output_path):
    command = [
        "ffmpeg",
        "-i", image_path,
        "-i", audio_path,
        "-filter_complex",
        "[1:a]aformat=channel_layouts=mono,showwaves=s=1200x400:mode=p2p:colors=blue@0.8[w];[0:v][w]overlay=(W-w)/2:(H-h)/2",
        "-c:v", "libx264",
        "-b:v", "2000k",
        "-c:a", "aac",
        "-b:a", "192k",
        "-y", output_path
    ]
    subprocess.run(command, check=True)

# Hàm xử lý sự kiện khi nhấn nút "Tạo Video"
def generate_video(prompt, language, audio_file_pth, normalize_text, use_llm, content_type):
    # Bước 1: Tạo audio nếu chưa có
    if not os.path.exists("output.wav"):
        audio_file, metrics_text = predict(prompt, language, audio_file_pth, normalize_text, use_llm, content_type)
        if not audio_file:
            return None, metrics_text
    else:
        audio_file = "output.wav"

    # Bước 2: Tạo mô tả ảnh
    image_description = generate_image_description(prompt)

    # Bước 3: Gọi API tạo ảnh
    try:
        image = txt2img(image_description, width=800, height=600)
        if isinstance(image, str):  # Nếu có lỗi từ API
            return None, image
        
        # Lưu ảnh vào thư mục
        image_path = os.path.join(SAVE_DIR, "generated_image.png")
        image.save(image_path)
    except Exception as e:
        return None, f"Error generating image: {str(e)}"

    # Bước 4: Tạo video từ ảnh và audio
    video_output_path = os.path.join(SAVE_DIR, "output_video.mp4")
    try:
        create_video(image_path, audio_file, video_output_path)
    except Exception as e:
        return None, f"Error creating video: {str(e)}"

    return video_output_path, "Video created successfully!"

# Thư mục lưu trữ ảnh và video
SAVE_DIR = "generated_images"
Path(SAVE_DIR).mkdir(exist_ok=True)

# Hàm dự đoán và tạo audio
@spaces.GPU
def predict(
    prompt,
    language,
    audio_file_pth,
    normalize_text=True,
    use_llm=False,
    content_type="Theo yêu cầu",
):
    if use_llm:
        print("I: Generating text with LLM...")
        generated_text = create_content(prompt, content_type, language)
        print(f"Generated text: {generated_text}")
        prompt = generated_text
    elif content_type in ["tin thể thao", "tin bóng đá"]:
        print("I: Fetching sports news...")
        news_type = "football" if content_type == "tin bóng đá" else "all"
        sports_content = get_sports_news_content(news_type, language, 5)
        print(f"Sports content: {sports_content}")
        prompt = sports_content
    if language not in supported_languages:
        metrics_text = gr.Warning(
            f"Language you put {language} in is not in our Supported Languages, please choose from dropdown"
        )
        return (None, metrics_text)
    speaker_wav = audio_file_pth
    if len(prompt) < 2:
        metrics_text = gr.Warning("Please give a longer prompt text")
        return (None, metrics_text)
    try:
        metrics_text = ""
        t_latent = time.time()
        try:
            (
                gpt_cond_latent,
                speaker_embedding,
            ) = MODEL.get_conditioning_latents(
                audio_path=speaker_wav,
                gpt_cond_len=30,
                gpt_cond_chunk_len=4,
                max_ref_length=60,
            )
        except Exception as e:
            print("Speaker encoding error", str(e))
            metrics_text = gr.Warning(
                "It appears something wrong with reference, did you unmute your microphone?"
            )
            return (None, metrics_text)
        prompt = re.sub("([^\x00-\x7F]|\w)(\.|\。|\?)", r"\1 \2\2", prompt)
        if normalize_text and language == "vi":
            prompt = normalize_vietnamese_text(prompt)
        print("I: Generating new audio...")
        t0 = time.time()
        out = MODEL.inference(
            prompt,
            language,
            gpt_cond_latent,
            speaker_embedding,
            repetition_penalty=5.0,
            temperature=0.75,
            enable_text_splitting=True,
        )
        inference_time = time.time() - t0
        print(f"I: Time to generate audio: {round(inference_time*1000)} milliseconds")
        metrics_text += (
            f"Time to generate audio: {round(inference_time*1000)} milliseconds\n"
        )
        real_time_factor = (time.time() - t0) / out["wav"].shape[-1] * 24000
        print(f"Real-time factor (RTF): {real_time_factor}")
        metrics_text += f"Real-time factor (RTF): {real_time_factor:.2f}\n"
        keep_len = calculate_keep_len(prompt, language)
        out["wav"] = out["wav"][:keep_len]
        torchaudio.save("output.wav", torch.tensor(out["wav"]).unsqueeze(0), 24000)
    except RuntimeError as e:
        if "device-side assert" in str(e):
            print(
                f"Exit due to: Unrecoverable exception caused by language:{language} prompt:{prompt}",
                flush=True,
            )
            gr.Warning("Unhandled Exception encounter, please retry in a minute")
            print("Cuda device-assert Runtime encountered need restart")
            error_time = datetime.datetime.now().strftime("%d-%m-%Y-%H:%M:%S")
            error_data = [
                error_time,
                prompt,
                language,
                audio_file_pth,
            ]
            error_data = [str(e) if type(e) != str else e for e in error_data]
            print(error_data)
            print(speaker_wav)
            write_io = StringIO()
            csv.writer(write_io).writerows([error_data])
            csv_upload = write_io.getvalue().encode()
            filename = error_time + "_" + str(uuid.uuid4()) + ".csv"
            print("Writing error csv")
            error_api = HfApi()
            error_api.upload_file(
                path_or_fileobj=csv_upload,
                path_in_repo=filename,
                repo_id="coqui/xtts-flagged-dataset",
                repo_type="dataset",
            )
            speaker_filename = error_time + "_reference_" + str(uuid.uuid4()) + ".wav"
            error_api = HfApi()
            error_api.upload_file(
                path_or_fileobj=speaker_wav,
                path_in_repo=speaker_filename,
                repo_id="coqui/xtts-flagged-dataset",
                repo_type="dataset",
            )
            space = api.get_space_runtime(repo_id=repo_id)
            if space.stage != "BUILDING":
                api.restart_space(repo_id=repo_id)
            else:
                print("TRIED TO RESTART but space is building")
        else:
            if "Failed to decode" in str(e):
                print("Speaker encoding error", str(e))
                metrics_text = gr.Warning(
                    "It appears something wrong with reference, did you unmute your microphone?"
                )
            else:
                print("RuntimeError: non device-side assert error:", str(e))
                metrics_text = gr.Warning(
                    "Something unexpected happened please retry again."
                )
            return (None, metrics_text)
    return ("output.wav", metrics_text)

# Hàm xử lý tin thể thao
def generate_sports_news(news_type, language, limit):
    """Tạo nội dung tin thể thao"""
    try:
        content = get_sports_news_content(news_type, language, limit)
        return content, "Đã tạo tin thể thao thành công!"
    except Exception as e:
        return f"Lỗi khi tạo tin thể thao: {str(e)}", "Có lỗi xảy ra"

# Giao diện Gradio
with gr.Blocks(analytics_enabled=False, title="TTS & Sports News Generator", theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # 🎙️ TTS & Sports News Generator ✨
        ### Tạo nội dung và chuyển đổi thành giọng nói - Powered by TDNM
        """
    )
    
    with gr.Tabs() as tabs:
        # Tab 1: TTS Generator
        with gr.Tab("🎙️ Tạo Giọng Nói", id="tts_tab"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown("### 📝 Nhập Nội Dung")
                    input_text_gr = gr.Textbox(
                        label="Bạn cần nội dung gì?",
                        info="Tôi có thể viết và thu âm luôn cho bạn",
                        value="Lời tự sự của AI, 150 từ",
                        lines=4,
                        placeholder="Nhập nội dung bạn muốn chuyển thành giọng nói..."
                    )
                    
                    with gr.Row():
                        language_gr = gr.Dropdown(
                            label="🌍 Ngôn ngữ",
                            choices=[
                                "vi", "en", "es", "fr", "de", "it", "pt", "pl", "tr", "ru", "nl", "cs", "ar", "zh-cn", "ja", "ko", "hu", "hi",
                            ],
                            value="vi",
                        )
                        content_type_dropdown = gr.Dropdown(
                            label="📋 Loại nội dung",
                            choices=["Theo yêu cầu", "triết lý sống", "tin thể thao", "tin bóng đá"],
                            value="Theo yêu cầu",
                        )
                    
                    with gr.Row():
                        normalize_text = gr.Checkbox(
                            label="Chuẩn hóa văn bản tiếng Việt",
                            value=True,
                        )
                        use_llm_checkbox = gr.Checkbox(
                            label="Sử dụng AI tạo nội dung",
                            value=True,
                        )
                    
                    ref_gr = gr.Audio(
                        label="🎤 Giọng mẫu (Reference Audio)",
                        type="filepath",
                        value="nam-tai-llieu.wav",
                    )
                    
                    with gr.Row():
                        tts_button = gr.Button(
                            "🗣️ Tạo Giọng Nói",
                            variant="primary",
                            size="lg"
                        )
                        video_button = gr.Button(
                            "🎥 Tạo Video", 
                            variant="secondary",
                            size="lg"
                        )
                
                with gr.Column(scale=1):
                    gr.Markdown("### 🎵 Kết Quả")
                    audio_gr = gr.Audio(label="Audio đã tạo", autoplay=True)
                    out_text_gr = gr.Textbox(label="📊 Thông tin", lines=3)
                    video_output = gr.Video(label="Video đã tạo")
                    video_status = gr.Textbox(label="Trạng thái video")
        
        # Tab 2: Sports News Generator
        with gr.Tab("⚽ Tin Thể Thao", id="sports_tab"):
            with gr.Row():
                with gr.Column(scale=1):
                    gr.Markdown(
                        """
                        ### 📰 Tạo Tin Thể Thao Tự Động
                        Lấy tin thể thao mới nhất từ các nguồn quốc tế và tạo nội dung bằng tiếng Việt
                        """
                    )
                    
                    with gr.Row():
                        sports_news_type = gr.Dropdown(
                            label="🏆 Loại tin thể thao",
                            choices=["all", "football"],
                            value="all",
                            info="Chọn loại tin thể thao bạn muốn"
                        )
                        sports_language = gr.Dropdown(
                            label="🌍 Ngôn ngữ nguồn",
                            choices=["vi", "en"],
                            value="vi",
                            info="Ngôn ngữ của nguồn tin"
                        )
                    
                    sports_limit = gr.Slider(
                        label="📊 Số lượng tin",
                        minimum=1,
                        maximum=10,
                        value=5,
                        step=1,
                        info="Số lượng tin thể thao muốn lấy"
                    )
                    
                    with gr.Row():
                        generate_sports_btn = gr.Button(
                            "📰 Tạo Tin Thể Thao",
                            variant="primary",
                            size="lg"
                        )
                        use_sports_for_tts_btn = gr.Button(
                            "🎙️ Chuyển Sang TTS",
                            variant="secondary",
                            size="lg"
                        )
                
                with gr.Column(scale=1):
                    gr.Markdown("### 📋 Nội Dung Tin Thể Thao")
                    sports_content_output = gr.Textbox(
                        label="Nội dung tin thể thao",
                        lines=15,
                        placeholder="Nội dung tin thể thao sẽ hiển thị ở đây...",
                        show_copy_button=True
                    )
                    sports_status = gr.Textbox(
                        label="Trạng thái",
                        lines=2
                    )
        
        # Tab 3: Hướng dẫn
        with gr.Tab("📖 Hướng Dẫn", id="guide_tab"):
            gr.Markdown(
                """
                ## 🚀 Hướng Dẫn Sử Dụng
                
                ### 🎙️ Tab Tạo Giọng Nói:
                1. **Nhập nội dung**: Gõ văn bản bạn muốn chuyển thành giọng nói
                2. **Chọn ngôn ngữ**: Hiện tại hỗ trợ tiếng Việt và nhiều ngôn ngữ khác
                3. **Chọn loại nội dung**: 
                   - "Theo yêu cầu": Sử dụng nội dung bạn nhập
                   - "triết lý sống": AI tạo nội dung về triết lý
                   - "tin thể thao": AI tạo tin thể thao
                   - "tin bóng đá": AI tạo tin bóng đá
                4. **Upload giọng mẫu**: Chọn file audio làm mẫu giọng nói
                5. **Nhấn "Tạo Giọng Nói"**: Hệ thống sẽ tạo audio
                6. **Tạo Video** (tùy chọn): Tạo video từ audio và hình ảnh
                
                ### ⚽ Tab Tin Thể Thao:
                1. **Chọn loại tin**: Tất cả tin thể thao hoặc chỉ bóng đá
                2. **Chọn ngôn ngữ nguồn**: Tiếng Việt hoặc tiếng Anh
                3. **Chọn số lượng tin**: Từ 1-10 tin
                4. **Nhấn "Tạo Tin Thể Thao"**: Lấy tin mới nhất
                5. **Chuyển sang TTS**: Sử dụng nội dung tin để tạo giọng nói
                
                ### 🔧 Tính Năng Nổi Bật:
                - ✅ TTS tiếng Việt chất lượng cao với viXTTS
                - ✅ Tạo nội dung tự động bằng AI
                - ✅ Lấy tin thể thao thời gian thực
                - ✅ Tạo video từ audio và hình ảnh
                - ✅ Hỗ trợ nhiều ngôn ngữ
                - ✅ Giao diện thân thiện, dễ sử dụng
                
                ### 📞 Hỗ Trợ:
                Nếu gặp vấn đề, vui lòng liên hệ: **https://www.tdn-m.com**
                """
            )
    
    # Event handlers cho tab TTS
    tts_button.click(
        predict,
        [
            input_text_gr,
            language_gr,
            ref_gr,
            normalize_text,
            use_llm_checkbox,
            content_type_dropdown,
        ],
        outputs=[audio_gr, out_text_gr],
        api_name="predict",
    )
    
    video_button.click(
        generate_video,
        inputs=[
            input_text_gr,
            language_gr,
            ref_gr,
            normalize_text,
            use_llm_checkbox,
            content_type_dropdown,
        ],
        outputs=[video_output, video_status],
    )
    
    # Event handlers cho tab tin thể thao
    generate_sports_btn.click(
        generate_sports_news,
        inputs=[
            sports_news_type,
            sports_language,
            sports_limit
        ],
        outputs=[sports_content_output, sports_status]
    )
    
    # Hàm chuyển nội dung tin thể thao sang tab TTS
    def transfer_sports_to_tts(sports_content):
        """Chuyển nội dung tin thể thao sang tab TTS"""
        if not sports_content or sports_content.strip() == "":
            return "Vui lòng tạo tin thể thao trước!", gr.update(selected="tts_tab")
        return sports_content, gr.update(selected="tts_tab")
    
    use_sports_for_tts_btn.click(
        transfer_sports_to_tts,
        inputs=[sports_content_output],
        outputs=[input_text_gr, tabs]
    )

# Khởi chạy ứng dụng
if __name__ == "__main__":
    demo.queue()
    demo.launch(
        debug=True, 
        show_api=True, 
        share=True,
        server_name="0.0.0.0",
        server_port=7860,
        favicon_path=None,
        ssl_verify=False
    )