import os
import tempfile
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from faster_whisper import WhisperModel
from pipeline.postprocess import process

app = FastAPI()

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

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

@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

    try:
        await websocket.send_json({"status": "connected"})
        print("Step 1: sent connected signal, 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": "transcribing"})
        print("Step 4: starting transcription")

        segments_raw, info = model.transcribe(
            tmp_path,
            language="en",
            word_timestamps=True,
            vad_filter=True,
        )
        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")

        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)
        try:
            await websocket.close()
        except Exception:
            pass