import os
import uuid
import torch
import yt_dlp
import whisperx
from fastapi import FastAPI, BackgroundTasks, Request
from fastapi.responses import HTMLResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from huggingface_hub import HfApi

# Enforce 50% CPU usage within PyTorch
torch.set_num_threads(1)

app = FastAPI()
app.mount("/static", StaticFiles(directory="static"), name="static")

# In-memory status tracker
jobs = {}

HF_TOKEN = os.getenv("HF_TOKEN")
DATASET_REPO = "your-username/whisper-transcripts" # Update with your dataset ID

def process_audio(job_id: str, url: str):
    jobs[job_id] = {"status": "downloading", "url": url}
    audio_path = f"/tmp/{job_id}.mp3"
    
    try:
        # 1. Extract and Download Audio
        ydl_opts = {'format': 'bestaudio/best', 'outtmpl': audio_path, 'quiet': True}
        with yt_dlp.YoutubeDL(ydl_opts) as ydl:
            ydl.download([url])
            
        jobs[job_id]["status"] = "transcribing"
        
        # 2. Transcription & Speaker Diarization (CPU forced to 1 thread)
        device = "cpu" 
        audio = whisperx.load_audio(audio_path)
        
        model = whisperx.load_model("base", device, compute_type="int8", threads=1)
        result = model.transcribe(audio)
        
        # Align timestamps and assign speakers
        model_a, metadata = whisperx.load_align_model(language_code=result["language"], device=device)
        result = whisperx.align(result["segments"], model_a, metadata, audio, device, return_char_alignments=False)
        
        diarize_model = whisperx.DiarizationPipeline(use_auth_token=HF_TOKEN, device=device)
        diarize_segments = diarize_model(audio)
        result = whisperx.assign_word_speakers(diarize_segments, result)
        
        # 3. Format to TXT
        output_txt = f"/tmp/{job_id}.txt"
        with open(output_txt, "w", encoding="utf-8") as f:
            for segment in result["segments"]:
                speaker = segment.get("speaker", "UNKNOWN")
                start = round(segment["start"], 2)
                end = round(segment["end"], 2)
                text = segment["text"]
                f.write(f"[{start} - {end}] {speaker}: {text}\n")
        
        jobs[job_id]["status"] = "pushing to dataset"
        
        # 4. Upload to Hugging Face Dataset for persistent download
        api = HfApi()
        api.upload_file(
            path_or_fileobj=output_txt,
            path_in_repo=f"{job_id}.txt",
            repo_id=DATASET_REPO,
            repo_type="dataset",
            token=HF_TOKEN
        )
        
        dataset_url = f"https://huggingface.co/datasets/{DATASET_REPO}/blob/main/{job_id}.txt"
        jobs[job_id] = {"status": "completed", "download_url": dataset_url}
        
    except Exception as e:
        jobs[job_id] = {"status": "error", "error": str(e)}
    finally:
        if os.path.exists(audio_path):
            os.remove(audio_path)

@app.get("/")
def read_index():
    with open("static/index.html") as f:
        return HTMLResponse(f.read())

@app.post("/start")
async def start_job(request: Request, background_tasks: BackgroundTasks):
    data = await request.json()
    url = data.get("url")
    job_id = str(uuid.uuid4())
    
    # Delegate the heavy lifting to the background queue
    background_tasks.add_task(process_audio, job_id, url)
    return JSONResponse({"job_id": job_id})

@app.get("/status/{job_id}")
def get_status(job_id: str):
    return JSONResponse(jobs.get(job_id, {"status": "not_found"}))