import os
import tempfile
import numpy as np
import torch
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel
from df.enhance import enhance, init_df, load_audio, save_audio
from pipeline.postprocess import process

app = FastAPI()

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

print("Loading Whisper model...", flush=True)
model = WhisperModel("large-v3", device="cpu", compute_type="int8")
print("Whisper model ready!")

print("Loading DeepFilterNet model...", flush=True)
df_model, df_state, _ = init_df()
print("DeepFilterNet model ready!")

USE_DEEPFILTERNET = False  # برای تست بدون DeepFilterNet، این رو موقتاً False کن

@app.get("/health")
async def health():
    return {"status": "ok", "message": "SubFlow backend is running"}


@app.websocket("/transcribe")
async def transcribe(websocket: WebSocket):
    await websocket.accept()
    tmp_path = None
    enhanced_path = None

    try:
        print("Step 1: waiting for bytes")
        audio_bytes = await websocket.receive_bytes()
        print(f"Step 2: received {len(audio_bytes)} bytes")
        await websocket.send_json({"status": "receiving"})

        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
            tmp.write(audio_bytes)
            tmp_path = tmp.name
        print(f"Step 3: saved to {tmp_path}")

        await websocket.send_json({"status": "cleaning"})
        print("Step 3.5: cleaning audio with DeepFilterNet")

        audio_to_transcribe = tmp_path

        if USE_DEEPFILTERNET:
            await websocket.send_json({"status": "cleaning"})
            print("Step 3.5: cleaning audio with DeepFilterNet")

            audio_data, _ = load_audio(tmp_path, sr=df_state.sr())
            enhanced_audio = enhance(df_model, df_state, audio_data, atten_lim_db=20)
            enhanced_path = tmp_path.replace(".wav", "_enhanced.wav")
            save_audio(enhanced_path, enhanced_audio, df_state.sr())
            audio_to_transcribe = enhanced_path
            print(f"Step 3.6: cleaned audio saved to {enhanced_path}")
        else:
            print("Step 3.5: skipping DeepFilterNet (disabled for A/B test)")

        await websocket.send_json({"status": "transcribing"})
        print("Step 4: starting transcription")

        segments_raw, info = model.transcribe(
            audio_to_transcribe,
            language="en",
            word_timestamps=True,
            vad_filter=True,
            vad_parameters=dict(
                min_silence_duration_ms=500,
                speech_pad_ms=200,
                threshold=0.1
            ),
        )

        print("Step 5: transcription done")

        await websocket.send_json({"status": "processing"})

        segments = []
        for seg in segments_raw:
            words = list(seg.words)
            if not words:
                continue
            segments.append({
                "text": seg.text.strip(),
                "start": words[0].start,
                "end": words[-1].end,
                "words": [{"word": w.word, "start": w.start, "end": w.end} for w in words],
            })

        print(f"Step 6: got {len(segments)} segments")
        for i, seg in enumerate(segments[:5]):
            print(f"  seg[{i}]: '{seg['text']}' | words={len(seg['words'])} | dur={seg['end']-seg['start']:.2f}s")

        srt_content = process(segments)
        print(f"Step 7: SRT built")

        await websocket.send_json({
            "status": "done",
            "srt": srt_content,
            "language": info.language,
            "duration": info.duration,
        })
        print("Step 8: done sent!")

    except WebSocketDisconnect:
        print("WebSocket disconnected by client")
    except Exception as e:
        print(f"ERROR: {type(e).__name__}: {e}")
        import traceback
        traceback.print_exc()
        try:
            await websocket.send_json({"status": "error", "message": str(e)})
        except Exception:
            pass
    finally:
        if tmp_path and os.path.exists(tmp_path):
            os.remove(tmp_path)
        if enhanced_path and os.path.exists(enhanced_path):
            os.remove(enhanced_path)
        try:
            await websocket.close()
        except Exception:
            pass