import json
import time
import uuid
import io
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from test_run import load_finetuned_model, transcribe_generator
import soundfile as sf
import uvicorn
import base64

# filepath: /home/anuran/s2t-model/server.py

app = FastAPI(title="Whisper Indic Voice Agent")

# Load model at startup
model = None
processor = None

LANGUAGE_CODE_MAP = {
  "hi": "hi-IN",
  "en": "en-US",
  "bn": "bn-IN",
  "ta": "ta-IN",
  "te": "te-IN",
  "mr": "mr-IN",
  "gu": "gu-IN",
  "kn": "kn-IN",
  "ml": "ml-IN",
  "pa": "pa-IN",
  "ur": "ur-IN",
  "or": "or-IN",
}


@app.on_event("startup")
async def startup_event():
  global model, processor
  print("Loading model...")
  model, processor = load_finetuned_model()
  print("Model loaded successfully.")


@app.websocket("/ws/transcribe")
async def websocket_transcribe(websocket: WebSocket):
  await websocket.accept()
  request_id_prefix = time.strftime("%Y%m%d")

  try:
    while True:
      # Receive JSON message from client
      print("Received WebSocket request.")
      message = await websocket.receive_text()
      payload = json.loads(message)

      # Extract and decode base64 audio data
      audio_info = payload.get("audio", {})
      b64_data = audio_info.get("data", "")
      encoding = audio_info.get("encoding", "audio/wav")
      sample_rate = int(audio_info.get("sample_rate", "16000"))

      raw_audio_bytes = base64.b64decode(b64_data)

      request_id = f"{request_id_prefix}_{uuid.uuid4()}"
      session_id = str(uuid.uuid4())

      # Send START_SPEECH event
      start_event = {
        "type": "events",
        "data": {
          "signal_type": "START_SPEECH",
          "occured_at": time.time(),
          "session_id": session_id,
        }
      }
      await websocket.send_text(json.dumps(start_event))

      # Wrap bytes in BytesIO
      audio_buffer = io.BytesIO(raw_audio_bytes)
      processing_start = time.time()

      detected_language = None
      transcription_text = None

      for result in transcribe_generator(audio_buffer, model, processor, language=None):
        if result["type"] == "language_detected":
          detected_language = result["language"]
        elif result["type"] == "transcription":
          transcription_text = result["transcription"]
          if detected_language is None:
            detected_language = result["language"]

      processing_latency = time.time() - processing_start

      print("Latency (in s) = ", processing_latency)

      # Send END_SPEECH event
      end_event = {
        "type": "events",
        "data": {
          "signal_type": "END_SPEECH",
          "occured_at": time.time(),
          "session_id": session_id,
        }
      }
      await websocket.send_text(json.dumps(end_event))

      # Compute audio duration from buffer
      audio_buffer.seek(0)
      try:
        audio_data, sr = sf.read(audio_buffer)
        audio_duration = len(audio_data) / sr
      except Exception:
        audio_duration = 0.0

      language_code = LANGUAGE_CODE_MAP.get(detected_language, f"{detected_language}-IN") if detected_language else "unknown"

      # Send transcription data
      data_event = {
        "type": "data",
        "data": {
          "request_id": request_id,
          "transcript": transcription_text or "",
          "timestamps": None,
          "diarized_transcript": None,
          "language_code": language_code,
          "language_probability": None,
          "metrics": {
            "audio_duration": round(audio_duration, 2),
            "processing_latency": processing_latency,
          }
        }
      }
      await websocket.send_text(json.dumps(data_event))

  except WebSocketDisconnect:
    print(f"Client disconnected")
  except Exception as e:
    print(f"Error: {e}")
    try:
      await websocket.close(code=1011, reason=str(e))
    except Exception:
      pass


@app.get("/health")
async def health():
  return {"status": "ok", "model_loaded": model is not None}


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