import os
import uuid
import shutil
import asyncio
import aiohttp
import librosa
import numpy as np
import soundfile as sf
import json
from fastapi import FastAPI, UploadFile, File, Form, Body
from fastapi.responses import FileResponse, JSONResponse
from fastapi.templating import Jinja2Templates
from fastapi.requests import Request
from pydantic import BaseModel
from typing import List, Optional

app = FastAPI()

# --- لیست کارگرها ---
WORKER_URLS_LIST = [
    "https://ezmary-taqviat-sadaworker1.hf.space",
    "https://ezmary-taqviat-sadaworker2.hf.space",
    "https://ezmary-taqviat-sadaworker3.hf.space",
    "https://ezmary-taqviat-sadaworker4.hf.space",
    "https://ezmary-taqviat-sadaworker5.hf.space",
    "https://eltafjan-taqviat-sadaworker6.hf.space",
    "https://eltafjan-taqviat-sadaworker7.hf.space",
    "https://eltafjan-taqviat-sadaworker8.hf.space",
    "https://eltafjan-taqviat-sadaworker9.hf.space",
    "https://eltafjan-taqviat-sadaworker10.hf.space",
    "https://eltafjan-taqviat-sadaworker11.hf.space",
    "https://eltafjan-taqviat-sadaworker12.hf.space",
    "https://eltafjan-taqviat-sadaworker13.hf.space",
    "https://eltafjan-taqviat-sadaworker14.hf.space",
    "https://eltafjan-taqviat-sadaworker15.hf.space",
    "https://eltafjan-taqviat-sadaworker16.hf.space",
    "https://eltafjan-taqviat-sadaworker17.hf.space",
    "https://eltafjan-taqviat-sadaworker18.hf.space",
    "https://eltafjan-taqviat-sadaworker19.hf.space",
    "https://eltafjan-taqviat-sadaworker20.hf.space"
]

os.makedirs("temp", exist_ok=True)
os.makedirs("results", exist_ok=True)
templates = Jinja2Templates(directory="templates")

# --- مدیریت کارگرها (چرخشی) ---
class AtomicWorkerManager:
    def __init__(self, urls):
        self.urls = urls
        self.total_workers = len(urls)
        self.current_index = 0
        self.lock = asyncio.Lock()

    async def get_next_worker(self):
        async with self.lock:
            url = self.urls[self.current_index]
            self.current_index = (self.current_index + 1) % self.total_workers
            return url

worker_manager = AtomicWorkerManager(WORKER_URLS_LIST)

# --- مدل‌های داده ---
class ChunkInfo(BaseModel):
    index: int
    worker_url: str
    task_id: str

class ProjectState(BaseModel):
    job_id: str
    total_chunks: int
    chunks: List[ChunkInfo]

# --- توابع کمکی ---
def find_split_points(audio_path, sr=24000):
    try:
        y, _ = librosa.load(audio_path, sr=sr)
    except:
        data, samplerate = sf.read(audio_path)
        if len(data.shape) > 1: data = np.mean(data, axis=1)
        if samplerate != sr:
            # Simple resampling if needed, ideally use librosa.resample
            # But here relying on librosa.load mostly
            pass 
        y = data
        
    total_samples = len(y)
    split_points = [0]
    current_pos = 0

    # منطق برش: بین 30 تا 60 ثانیه
    min_duration = 30 * sr
    max_duration = 60 * sr

    while current_pos < total_samples:
        # هدف نهایی (60 ثانیه بعد)
        target = current_pos + max_duration 
        
        # اگر به انتهای فایل نزدیکیم
        if target >= total_samples:
            split_points.append(total_samples)
            break
        
        # بازه جستجو برای سکوت: از ثانیه 30 تا 60
        search_start = current_pos + min_duration
        search_end = target # تا خود 60 ثانیه
        
        if search_start >= total_samples:
             split_points.append(total_samples)
             break

        region = y[search_start:search_end]
        
        # پیدا کردن کمترین انرژی (سکوت)
        if len(region) > 0:
            rms = librosa.feature.rms(y=region, frame_length=1024, hop_length=512)[0]
            min_idx = np.argmin(rms)
            # تبدیل ایندکس فریم به ایندکس سمپل
            cut_point = search_start + (min_idx * 512)
        else:
            cut_point = target

        split_points.append(cut_point)
        current_pos = cut_point
        
    return split_points, y

async def submit_to_worker(session, worker_url, chunk_path, params):
    try:
        with open(chunk_path, 'rb') as f_c:
            data = aiohttp.FormData()
            data.add_field('audio_file', f_c, filename='input.wav', content_type='audio/wav')
            # ارسال تنظیمات
            data.add_field('solver', params.get('solver', 'Midpoint'))
            data.add_field('nfe', str(params.get('nfe', 64)))
            data.add_field('tau', str(params.get('tau', 0.5)))
            data.add_field('denoising', str(params.get('denoising', 'false')))

            async with session.post(f"{worker_url}/submit", data=data, timeout=120) as resp:
                if resp.status == 200:
                    js = await resp.json()
                    return js.get("task_id")
    except Exception as e:
        print(f"Submit error to {worker_url}: {e}")
        return None
    return None

async def check_worker_status(session, worker_url, task_id):
    try:
        async with session.get(f"{worker_url}/result/{task_id}") as resp:
            if resp.status == 200:
                return "completed", await resp.read()
            elif resp.status == 202:
                return "processing", None 
            elif resp.status == 404:
                return "processing", None 
            else:
                return "failed", None
    except:
        return "processing", None

@app.get("/")
def home(request: Request):
    return templates.TemplateResponse("index.html", {"request": request})

@app.post("/upload")
async def start_process(
    source_audio: UploadFile = File(...),
    solver: str = Form("Midpoint"),
    nfe: int = Form(64),
    tau: float = Form(0.5),
    denoising: bool = Form(False)
):
    job_id = str(uuid.uuid4())
    os.makedirs(f"temp/{job_id}", exist_ok=True)
    
    src_path = f"temp/{job_id}/src.wav"
    with open(src_path, "wb") as b: shutil.copyfileobj(source_audio.file, b)
    
    # تنظیمات برای ارسال به کارگر
    worker_params = {
        "solver": solver,
        "nfe": nfe,
        "tau": tau,
        "denoising": "true" if denoising else "false"
    }
    
    sr = 44100 # Resemble معمولا با 44100 کار میکند
    
    # برش فایل
    split_points, y = find_split_points(src_path, sr)
    total_chunks = len(split_points) - 1
    
    chunks_metadata = []
    
    async with aiohttp.ClientSession() as session:
        tasks = []
        for i in range(total_chunks):
            start = split_points[i]
            end = split_points[i+1]
            chunk_audio = y[start:end]
            
            # رد کردن تکه‌های خیلی کوتاه (زیر 0.5 ثانیه)
            if len(chunk_audio) < 0.5 * sr:
                chunks_metadata.append({"index": i, "worker_url": "skip", "task_id": "skip"})
                continue

            chunk_path = f"temp/{job_id}/chunk_{i}.wav"
            sf.write(chunk_path, chunk_audio, sr)
            
            worker_url = await worker_manager.get_next_worker()
            
            tasks.append(submit_to_worker(session, worker_url, chunk_path, worker_params))
            
            chunks_metadata.append({
                "index": i,
                "worker_url": worker_url,
                "task_id": "pending" 
            })
            
        results = await asyncio.gather(*tasks)
        
        active_task_idx = 0
        for i in range(len(chunks_metadata)):
            if chunks_metadata[i]["task_id"] == "skip": continue
            
            task_id = results[active_task_idx]
            active_task_idx += 1
            
            if task_id:
                chunks_metadata[i]["task_id"] = task_id
            else:
                chunks_metadata[i]["task_id"] = "failed"
    
    return {
        "job_id": job_id,
        "total_chunks": total_chunks,
        "chunks": chunks_metadata,
        "status": "started"
    }

@app.post("/check_status")
async def check_status(project: ProjectState):
    final_filename = f"enhanced_{project.job_id}.wav"
    if os.path.exists(f"results/{final_filename}"):
        return {"status": "completed", "progress": 100, "filename": final_filename}
    
    completed_count = 0
    audio_parts = {}
    failed_any = False
    
    async with aiohttp.ClientSession() as session:
        tasks = []
        task_indices = []
        
        for chunk in project.chunks:
            if chunk.task_id == "skip":
                completed_count += 1
                audio_parts[chunk.index] = None 
                continue
            if chunk.task_id == "failed":
                failed_any = True
                continue
                
            tasks.append(check_worker_status(session, chunk.worker_url, chunk.task_id))
            task_indices.append(chunk.index)
            
        results = await asyncio.gather(*tasks)
        
        for i, (status, data) in enumerate(results):
            idx = task_indices[i]
            if status == "completed" and data:
                completed_count += 1
                audio_parts[idx] = data
            elif status == "failed":
                failed_any = True
            
    progress = int((completed_count / project.total_chunks) * 100)
    
    if completed_count == project.total_chunks or (completed_count > 0 and progress > 98 and failed_any):
        try:
            full_audio = []
            # استفاده از سمپل ریت خروجی Resemble (معمولا 44100)
            target_sr = 44100 
            
            for i in range(project.total_chunks):
                if i in audio_parts:
                    if isinstance(audio_parts[i], bytes):
                        tmp_path = f"temp/part_{project.job_id}_{i}.wav"
                        with open(tmp_path, "wb") as f: f.write(audio_parts[i])
                        y, sr = librosa.load(tmp_path, sr=None) # Load native SR
                        full_audio.append(y)
                        target_sr = sr # Update SR based on worker output
                        os.remove(tmp_path)
                    else:
                        pass # Skip part (silence?)
                else:
                    # Missing part due to error
                    full_audio.append(np.zeros(int(target_sr * 1.0))) # 1 sec silence
            
            final_wav = np.concatenate(full_audio)
            sf.write(f"results/{final_filename}", final_wav, target_sr)
            
            shutil.rmtree(f"temp/{project.job_id}", ignore_errors=True)
            
            return {"status": "completed", "progress": 100, "filename": final_filename}
            
        except Exception as e:
            print(f"Stitch error: {e}")
            return {"status": "error", "progress": progress, "detail": str(e)}
    else:
        return {"status": "processing", "progress": progress}

@app.get("/download/{filename}")
def download_file(filename: str):
    path = f"results/{filename}"
    if os.path.exists(path):
        return FileResponse(path, filename=filename, media_type="audio/wav")
    return {"error": "File not found"}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)