import os
import io
import time
import logging
import torch
import torchaudio
from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from transformers import pipeline
from huggingface_hub import login

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('transcription.log')
    ]
)
logger = logging.getLogger(__name__)

# Set HF_HOME to a writable directory
os.environ['HF_HOME'] = '/app/.cache/huggingface'

# Handle authentication if token is provided
hf_token = os.getenv("HF_TOKEN")
if hf_token:
    logger.info("Logging in to Hugging Face with provided token")
    login(token=hf_token)

# Initialize FastAPI app
app = FastAPI()

# Enable CORS for all origins (adjust for production)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Load the pipeline with device set to GPU if available
MODEL_NAME = "STT-Darija-ORG/wav2vec2-xlsr-300m-darija-augmented"
device = 0 if torch.cuda.is_available() else -1
logger.info(f"Loading model {MODEL_NAME} on device: {'GPU' if device == 0 else 'CPU'}")
start_time = time.time()
pipe = pipeline("automatic-speech-recognition", model=MODEL_NAME, device=device)
logger.info(f"Model loaded in {time.time() - start_time:.2f} seconds")

@app.post("/transcribe")
async def transcribe_audio(file: UploadFile = File(...)):
    start_time = time.time()
    try:
        # Read the file content
        logger.info(f"Received file: {file.filename}, size: {file.size} bytes")
        content = await file.read()
        if len(content) == 0:
            logger.error("Uploaded file is empty")
            raise HTTPException(status_code=400, detail="Uploaded file is empty")

        # Load audio from bytes
        load_start = time.time()
        file_like = io.BytesIO(content)
        try:
            waveform, sample_rate = torchaudio.load(file_like)
            logger.info(f"Audio loaded, sample rate: {sample_rate} Hz, channels: {waveform.shape[0]}, duration: {waveform.shape[1]/sample_rate:.2f} seconds")
        except Exception as e:
            logger.error(f"Failed to load audio: {str(e)}")
            raise HTTPException(status_code=400, detail=f"Invalid audio file: {str(e)}")
        load_time = time.time() - load_start
        logger.info(f"Audio loading took {load_time:.2f} seconds")

        # Ensure mono audio
        if waveform.shape[0] > 1:
            logger.info("Converting multi-channel audio to mono")
            waveform = waveform[0:1, :]  # Take the first channel

        # Resample if necessary
        if sample_rate != 16000:
            logger.info(f"Resampling audio from {sample_rate} Hz to 16000 Hz")
            resample_start = time.time()
            transform = torchaudio.transforms.Resample(sample_rate, 16000)
            waveform = transform(waveform)
            logger.info(f"Resampling took {time.time() - resample_start:.2f} seconds")

        # Pass to pipeline
        logger.info("Starting transcription")
        transcribe_start = time.time()
        audio_np = waveform.numpy().flatten()
        result = pipe(audio_np)
        transcription = result["text"]
        transcribe_time = time.time() - transcribe_start
        logger.info(f"Transcription completed in {transcribe_time:.2f} seconds, result: {transcription}")

        total_time = time.time() - start_time
        logger.info(f"Total request processing time: {total_time:.2f} seconds")

        return {"transcription": transcription, "processing_time_seconds": round(total_time, 2)}
    except HTTPException as e:
        logger.error(f"HTTP error: {str(e)}")
        raise e
    except Exception as e:
        logger.error(f"Unexpected error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Transcription failed: {str(e)}")

@app.get("/health")
async def health_check():
    logger.info("Health check requested")
    return {"status": "healthy", "model": MODEL_NAME}

@app.get("/")
async def home():
    logger.info("Home endpoint accessed")
    return {
        "message": "Speech Recognition API",
        "endpoints": {
            "/transcribe": "POST - Upload audio file for transcription",
            "/health": "GET - Check API health"
        }
    }