import modal
import os

# --- Helper function to pre-cache Whisper model during container build ---
def download_whisper():
    from transformers import pipeline
    pipeline("automatic-speech-recognition", model="openai/whisper-tiny")

# --- Setup Modal Image with ML dependencies ---
image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install("ffmpeg")
    .pip_install(
        "torch", 
        "transformers", 
        "accelerate",
        "fastapi[standard]",
        "pydantic",
        "numpy"
    )
    .run_function(download_whisper)
)

app = modal.App("memrl-canvas-backend")

# --- Whisper ASR CPU-bound Class ---
@app.cls(cpu=1.0, image=image, timeout=300)
class WhisperASR:
    @modal.enter()
    def setup(self):
        from transformers import pipeline
        self.asr_pipe = pipeline(
            "automatic-speech-recognition",
            model="openai/whisper-tiny",
            device="cpu"
        )
        print("Whisper Tiny ASR loaded successfully on CPU.")

    @modal.method()
    def transcribe(self, audio_bytes: bytes) -> str:
        import tempfile
        import os
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
            f.write(audio_bytes)
            temp_path = f.name
            
        try:
            result = self.asr_pipe(temp_path)
            text = result.get("text", "").strip()
            print(f"ASR Transcribed: '{text}'")
        except Exception as e:
            print(f"ASR transcription failed: {str(e)}")
            text = ""
        finally:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        return text

# --- Gemma NLU GPU-bound Class ---
@app.cls(gpu="A10G", secrets=[modal.Secret.from_name("huggingface-secret")], image=image, timeout=600)
class GemmaModel:
    @modal.enter()
    def setup(self):
        import torch
        from transformers import AutoTokenizer, AutoModelForCausalLM
        
        model_id = "google/gemma-4-E2B-it"
        hf_token = os.environ.get("HF_TOKEN")
        if not hf_token:
            raise ValueError("HF_TOKEN was not found in Modal environment variables.")
            
        print("Loading Gemma 4 E2B IT Model on GPU...")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id,
            torch_dtype=torch.bfloat16,
            device_map="cuda",
            token=hf_token
        )
        print("Gemma Model loaded successfully on GPU VRAM.")

    @modal.method()
    def generate(self, prompt: str) -> str:
        import torch
        
        inputs = self.tokenizer(prompt, return_tensors="pt").to("cuda")
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_new_tokens=180,
                do_sample=False
            )
        response = self.tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True)
        return response

# --- FastAPI ASGI Application ---
from fastapi import FastAPI, HTTPException, Body

web_app = FastAPI(title="MemRL Canvas Backend API")

@web_app.post("/transcribe")
def api_transcribe(payload: dict = Body(...)):
    import base64
    
    expected_key = os.environ.get("MODAL_API_KEY")
    if expected_key and payload.get("api_key") != expected_key:
        raise HTTPException(status_code=401, detail="Unauthorized API Key.")
        
    audio_base64 = payload.get("audio_base64")
    if not audio_base64:
        raise HTTPException(status_code=400, detail="Missing audio_base64")
        
    try:
        audio_bytes = base64.b64decode(audio_base64)
    except Exception as e:
        raise HTTPException(status_code=400, detail=f"Invalid base64 encoding: {str(e)}")
        
    asr = WhisperASR()
    text = asr.transcribe.remote(audio_bytes)
    return {"text": text}

@web_app.post("/gemma")
def api_gemma(payload: dict = Body(...)):
    expected_key = os.environ.get("MODAL_API_KEY")
    if expected_key and payload.get("api_key") != expected_key:
        raise HTTPException(status_code=401, detail="Unauthorized API Key.")
        
    prompt = payload.get("prompt")
    if not prompt:
        raise HTTPException(status_code=400, detail="Missing prompt")
        
    gemma_model = GemmaModel()
    response = gemma_model.generate.remote(prompt)
    return {"response": response}

@app.function(image=image, secrets=[modal.Secret.from_name("huggingface-secret")])
@modal.asgi_app()
def api():
    return web_app