import os
import asyncio
import uuid
import wave
import re
from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Request
from fastapi.responses import HTMLResponse, FileResponse
from fastapi.staticfiles import StaticFiles
from google import genai
from google.genai import types

# --- مدیریت کلیدهای API ---
class ApiKeyManager:
    def __init__(self, api_keys_str: str):
        if not api_keys_str:
            raise ValueError("متغیر ALL_GEMINI_API_KEYS پیدا نشد یا خالی است!")
        self.keys = [key.strip() for key in api_keys_str.split(',') if key.strip()]
        if not self.keys:
            raise ValueError("هیچ کلید معتبری در متغیر ALL_GEMINI_API_KEYS یافت نشد.")
        print(f"تعداد {len(self.keys)} کلید API با موفقیت بارگذاری شد.")
        self._index = 0
        self._lock = asyncio.Lock()

    async def get_next_key(self) -> tuple[int, str]:
        async with self._lock:
            key_index = self._index
            api_key = self.keys[key_index]
            self._index = (self._index + 1) % len(self.keys)
            return key_index, api_key

ALL_API_KEYS = os.environ.get("ALL_GEMINI_API_KEYS")
api_key_manager = ApiKeyManager(ALL_API_KEYS)

# --- تنظیمات عمومی Gemini ---
MODEL = "models/gemini-2.5-flash-native-audio-preview-09-2025"
CONFIG = types.LiveConnectConfig(
    response_modalities=["AUDIO"],
    speech_config=types.SpeechConfig(
        voice_config=types.VoiceConfig(
            prebuilt_voice_config=types.PrebuiltVoiceConfig(voice_name="Zephyr")
        )
    ),
)

app = FastAPI()

TEMP_DIR = "/tmp/temp_audio"
os.makedirs(TEMP_DIR, exist_ok=True)
app.mount("/audio", StaticFiles(directory=TEMP_DIR), name="audio")

# --- تنظیمات فایل صوتی ---
SAMPLE_RATE = 24000
CHANNELS = 1
SAMPLE_WIDTH = 2
AUDIO_BUFFER_SIZE = 8192

def split_text_into_chunks(text: str, max_chunk_size: int = 800):
    if len(text) <= max_chunk_size:
        return [text]
    sentences = re.split(r'(?<=[.!?؟])\s+', text.strip())
    chunks, current_chunk = [], ""
    for sentence in sentences:
        if not sentence: continue
        if len(current_chunk) + len(sentence) + 1 > max_chunk_size and current_chunk:
            chunks.append(current_chunk.strip())
            current_chunk = sentence
        else:
            current_chunk += (" " + sentence) if current_chunk else sentence
    if current_chunk: chunks.append(current_chunk.strip())
    return chunks

async def process_and_stream_one_chunk(websocket: WebSocket, chunk: str, request_id: str, chunk_num: int, total_chunks: int):
    max_retries = len(api_key_manager.keys)
    log_prefix = f"[{request_id}][قطعه {chunk_num}/{total_chunks}]"
    
    for attempt in range(max_retries):
        key_index, api_key = await api_key_manager.get_next_key()
        
        unique_id_for_chunk = str(uuid.uuid4())[:8]
        tts_prompt = f"Please ignore the ID in brackets and just read the text: '{chunk}' [ID: {unique_id_for_chunk}]"

        print(f"{log_prefix} تلاش {attempt + 1}/{max_retries} با کلید {key_index}...")

        try:
            client = genai.Client(http_options={"api_version": "v1beta"}, api_key=api_key)
            audio_buffer = bytearray()
            full_chunk_audio = []
            
            async with client.aio.live.connect(model=MODEL, config=CONFIG) as session:
                await session.send(input=tts_prompt, end_of_turn=True)
                turn = session.receive()
                
                first_response = await asyncio.wait_for(anext(turn), timeout=10.0)
                
                print(f"{log_prefix} ارتباط موفق با کلید {key_index}. شروع استریم...")
                
                if data := first_response.data:
                    full_chunk_audio.append(data)
                    audio_buffer.extend(data)

                async for response in turn:
                    if data := response.data:
                        full_chunk_audio.append(data)
                        audio_buffer.extend(data)
                        if len(audio_buffer) >= AUDIO_BUFFER_SIZE:
                            await websocket.send_bytes(audio_buffer)
                            audio_buffer = bytearray()
            
            if audio_buffer:
                await websocket.send_bytes(audio_buffer)
            
            return b"".join(full_chunk_audio)

        except Exception as e:
            print(f"{log_prefix} خطا با کلید {key_index}: {e}. تلاش با کلید بعدی...")
    
    raise Exception(f"پردازش {log_prefix} با تمام {max_retries} کلید API ناموفق بود.")


@app.get("/", response_class=HTMLResponse)
async def read_root():
    # بررسی می‌کنیم که فایل index.html کجا قرار دارد و همان را ارسال می‌کنیم
    if os.path.exists("index.html"):
        return FileResponse("index.html")
    elif os.path.exists("templates/index.html"):
        return FileResponse("templates/index.html")
    return HTMLResponse("<h1>خطا: فایل index.html پیدا نشد!</h1>", status_code=404)

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await websocket.accept()
    client_id = str(uuid.uuid4())[:6]
    print(f"کلاینت جدید [{client_id}] متصل شد.")

    try:
        while True:
            text_prompt_from_user = await websocket.receive_text()
            request_id = f"{client_id}-{str(uuid.uuid4())[:6]}"
            print(f"[{request_id}] درخواست جدید دریافت شد.")

            try:
                text_chunks = split_text_into_chunks(text_prompt_from_user)
                full_audio_for_file = []

                for i, chunk in enumerate(text_chunks):
                    chunk_audio = await process_and_stream_one_chunk(
                        websocket, chunk, request_id, i + 1, len(text_chunks)
                    )
                    full_audio_for_file.append(chunk_audio)
                
                filename = f"{request_id}.wav"
                filepath = os.path.join(TEMP_DIR, filename)
                with wave.open(filepath, 'wb') as wf:
                    wf.setnchannels(CHANNELS); wf.setsampwidth(SAMPLE_WIDTH)
                    wf.setframerate(SAMPLE_RATE); wf.writeframes(b"".join(full_audio_for_file))
                print(f"[{request_id}] فایل صوتی کامل در {filepath} ذخیره شد.")
                await websocket.send_json({"event": "STREAM_ENDED", "url": f"/audio/{filename}"})

                print(f"[{request_id}] استریم با موفقیت به پایان رسید.")

            except Exception as e:
                error_message = f"خطای بحرانی: {e}"
                print(f"[{request_id}] {error_message}")
                await websocket.send_json({"event": "ERROR", "message": str(e)})

    except WebSocketDisconnect:
        print(f"کلاینت [{client_id}] اتصال را قطع کرد.")
    except Exception as e:
        print(f"[{client_id}] یک خطای پیش‌بینی‌نشده در WebSocket رخ داد: {e}")
    finally:
        print(f"ارتباط با کلاینت [{client_id}] بسته شد.")