import os
import soundfile as sf
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from pydantic import BaseModel, Field
import uvicorn
import redis
import hashlib
import json
from kafka import KafkaProducer, KafkaConsumer
import threading
import time
from tools.i18n.i18n import I18nAuto
from GPT_SoVITS.inference_webui import change_gpt_weights, change_sovits_weights, get_tts_wav
from dotenv import load_dotenv
import os
import torch

# 加载 .env 文件
load_dotenv()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# FastAPI configuration
app = FastAPI()
i18n = I18nAuto()

# CORS configuration
ALLOWED_ORIGINS = os.getenv('ALLOWED_ORIGINS').split(',')

# Redis configuration
REDIS_HOST = os.getenv('REDIS_HOST')
REDIS_PORT = int(os.getenv('REDIS_PORT'))
REDIS_DB = int(os.getenv('REDIS_TTS_DB'))
REDIS_PASSWORD = os.getenv('REDIS_PASSWORD')

# Kafka configuration
KAFKA_BROKER = os.getenv('KAFKA_BROKER')
KAFKA_TOPIC = os.getenv('KAFKA_TTS_TOPIC')
# KAFKA_GROUP_ID = 'tts_group'
KAFKA_CONSUMER_THREADS = 1

# TTS configuration
GPT_MODEL_PATH = os.getenv('GPT_MODEL_PATH')
SOVITS_MODEL_PATH = os.getenv('SOVITS_MODEL_PATH')
REF_AUDIO_PATH = os.getenv('REF_AUDIO_KO_PATH')
REF_TEXT_PATH = os.getenv('REF_TEXT_KO_PATH')
REF_LANGUAGE = os.getenv('REF_KO_LANGUAGE')
TARGET_LANGUAGE = os.getenv('TARGET_LANGUAGE')
OUTPUT_PATH = os.getenv('OUTPUT_PATH')

# Initialize FastAPI CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=ALLOWED_ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Initialize Redis client
redis_client = redis.Redis(
    host=REDIS_HOST, 
    port=REDIS_PORT, 
    db=REDIS_DB,
    password=REDIS_PASSWORD
)

# Initialize Kafka producer
kafka_producer = KafkaProducer(bootstrap_servers=KAFKA_BROKER)

class TTSRequest(BaseModel):
    text: str = Field(..., alias="text")

def get_audio_hash(text):
    return hashlib.md5(text.encode()).hexdigest()

# Initialize models at startup
print("Initializing models...")
change_gpt_weights(gpt_path=GPT_MODEL_PATH)
change_sovits_weights(sovits_path=SOVITS_MODEL_PATH)

# Read reference text
with open(REF_TEXT_PATH, 'r', encoding='utf-8') as file:
    ref_text = file.read()

print("Models initialized successfully.")

def synthesize(target_text, output_path):
    # Synthesize audio
    with torch.cuda.device(device):
        synthesis_result = get_tts_wav(ref_wav_path=REF_AUDIO_PATH, 
                                    prompt_text=ref_text, 
                                    prompt_language=i18n(REF_LANGUAGE), 
                                    text=target_text, 
                                    text_language=i18n(TARGET_LANGUAGE), top_p=1, temperature=1)
        
        result_list = list(synthesis_result)

        if result_list:
            last_sampling_rate, last_audio_data = result_list[-1]
            audio_hash = get_audio_hash(target_text)
            output_wav_path = os.path.join(output_path, f"{audio_hash}.wav")
            sf.write(output_wav_path, last_audio_data, last_sampling_rate)
            return output_wav_path
        else:
            return None

@app.post("/tts_ko")
async def synthesize_audio(request: TTSRequest):
    try:
        print(f"Received TTS request: {request.dict()}")
        target_text = request.text
        audio_hash = get_audio_hash(target_text)
        
        # Check Redis cache
        cached_audio = redis_client.get(audio_hash)
        if cached_audio:
            audio_info = json.loads(cached_audio)
            return FileResponse(audio_info['path'], media_type="audio/wav")

        # Check file system
        file_path = os.path.join(OUTPUT_PATH, f"{audio_hash}.wav")
        if os.path.exists(file_path):
            # Cache the file path in Redis
            redis_client.set(audio_hash, json.dumps({"path": file_path}))
            return FileResponse(file_path, media_type="audio/wav")

        # Send message to Kafka
        kafka_producer.send(KAFKA_TOPIC, json.dumps({
            'text': target_text,
            'audio_hash': audio_hash
        }).encode('utf-8'))

        # Wait for the audio to be generated (you might want to implement a more sophisticated waiting mechanism)
        for _ in range(60):  # Wait for up to 30 seconds
            if os.path.exists(file_path):
                return FileResponse(file_path, media_type="audio/wav")
            time.sleep(1)

        # If audio is not generated within the timeout
        raise HTTPException(status_code=504, detail="Audio generation timed out")
    except Exception as e:
        print(f"Error processing TTS request: {str(e)}")
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "TTS API is running"}

def kafka_consumer_thread():
    consumer = KafkaConsumer(
        KAFKA_TOPIC,
        bootstrap_servers=KAFKA_BROKER,
        # group_id=KAFKA_GROUP_ID,
        auto_offset_reset='latest',
        value_deserializer=lambda m: json.loads(m.decode('utf-8'))
    )

    for message in consumer:
        target_text = message.value['text']
        audio_hash = message.value['audio_hash']

        output_path = synthesize(target_text, OUTPUT_PATH)

        if output_path:
            redis_client.set(audio_hash, json.dumps({"path": output_path}))
            print(f"Audio synthesized successfully: {output_path}")
        else:
            print("Failed to synthesize audio")

if __name__ == "__main__":
    # Start Kafka consumer threads
    torch.cuda.set_device(device)
    for _ in range(KAFKA_CONSUMER_THREADS):
        consumer_thread = threading.Thread(target=kafka_consumer_thread)
        consumer_thread.start()

    uvicorn.run(app, host="0.0.0.0", port=6003)