import logging
import math
import time
import base64
import os
import subprocess
import tempfile
from typing import Dict, Any, Union, Tuple
from functools import wraps

from fastapi import FastAPI, Depends, HTTPException, File, UploadFile, Form, Header
from fastapi.encoders import jsonable_encoder
from pydantic import BaseModel
import jax.numpy as jnp
import numpy as np
from transformers.pipelines.audio_utils import ffmpeg_read
from whisper_jax import FlaxWhisperPipline

app = FastAPI(title="Whisper JAX: The Fastest Whisper API ⚡️")

logger = logging.getLogger("whisper-jax-app")
logger.setLevel(logging.DEBUG)
ch = logging.StreamHandler()
ch.setLevel(logging.DEBUG)
formatter = logging.Formatter("%(asctime)s;%(levelname)s;%(message)s", "%Y-%m-%d %H:%M:%S")
ch.setFormatter(formatter)
logger.addHandler(ch)

checkpoint = "openai/whisper-large-v3"

BATCH_SIZE = 32
CHUNK_LENGTH_S = 30
NUM_PROC = 32
FILE_LIMIT_MB = 10000

pipeline = FlaxWhisperPipline(checkpoint, dtype=jnp.bfloat16, batch_size=BATCH_SIZE)
stride_length_s = CHUNK_LENGTH_S / 6
chunk_len = round(CHUNK_LENGTH_S * pipeline.feature_extractor.sampling_rate)
stride_left = stride_right = round(stride_length_s * pipeline.feature_extractor.sampling_rate)
step = chunk_len - stride_left - stride_right

# Pre-compile step
logger.debug("Compiling forward call...")
start = time.time()
random_inputs = {
    "input_features": np.ones(
        (BATCH_SIZE, pipeline.model.config.num_mel_bins, 2 * pipeline.model.config.max_source_positions)
    )
}
random_timestamps = pipeline.forward(random_inputs, batch_size=BATCH_SIZE, return_timestamps=True)
compile_time = time.time() - start
logger.debug(f"Compiled in {compile_time}s")

class TranscribeAudioRequest(BaseModel):
    audio_base64: str
    task: str = "transcribe"
    return_timestamps: bool = False

def timeit(func):
    @wraps(func)
    async def wrapper(*args, **kwargs):
        start_time = time.time()
        result = await func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        if isinstance(result, dict):
            result['total_execution_time'] = execution_time
        else:
            result = {'result': result, 'total_execution_time': execution_time}
        return result
    return wrapper

def check_api_key(x_api_key: str = Header(...)):
    api_key = os.environ.get("WHISPER_API_KEY")
    if not api_key or x_api_key != api_key:
        raise HTTPException(status_code=401, detail="Invalid or missing API key")
    return x_api_key

def extract_audio_from_video(video_data: bytes) -> bytes:
    """Extract audio from video file using ffmpeg."""
    with tempfile.NamedTemporaryFile(suffix='.mp4', delete=True) as video_file:
        with tempfile.NamedTemporaryFile(suffix='.wav', delete=True) as audio_file:
            # Write video data to temporary file
            video_file.write(video_data)
            video_file.flush()
            
            try:
                # Extract audio to WAV format
                subprocess.run([
                    'ffmpeg',
                    '-i', video_file.name,
                    '-vn',  # Disable video
                    '-acodec', 'pcm_s16le',  # Convert to PCM WAV
                    '-ar', '16000',  # Set sample rate to 16kHz
                    '-ac', '1',  # Convert to mono
                    '-y',  # Overwrite output file
                    audio_file.name
                ], check=True, capture_output=True)
                
                # Read the extracted audio
                return audio_file.read()
                
            except subprocess.CalledProcessError as e:
                logger.error(f"FFmpeg error: {e.stderr.decode() if e.stderr else str(e)}")
                raise HTTPException(
                    status_code=400,
                    detail="Error extracting audio from video file. Make sure it's a valid video file."
                )

def is_video_file(file_name: str) -> bool:
    """Check if the file is a video based on its extension."""
    video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv'}
    return any(file_name.lower().endswith(ext) for ext in video_extensions)

@app.post("/transcribe_audio_file")
@timeit
async def transcribe_audio_file(
    file: UploadFile = File(...),
    task: str = Form("transcribe"),
    return_timestamps: bool = Form(False),
    api_key: str = Depends(check_api_key)
) -> Dict[str, Any]:
    logger.debug("Starting transcribe_audio_file function")
    logger.debug(f"Received parameters - task: {task}, return_timestamps: {return_timestamps}")
    
    try:
        file_data = await file.read()
        file_size = len(file_data)
        file_size_mb = file_size / (1024 * 1024)
        logger.debug(f"File size: {file_size} bytes ({file_size_mb:.2f}MB)")

        # Check if the file is a video and extract audio if needed
        if is_video_file(file.filename):
            logger.debug("Processing video file")
            try:
                file_data = extract_audio_from_video(file_data)
                logger.debug("Successfully extracted audio from video")
            except Exception as e:
                logger.error(f"Error processing video file: {str(e)}", exc_info=True)
                raise HTTPException(
                    status_code=500,
                    detail=f"Error processing video file: {str(e)}"
                )

        return await process_audio(file_data, file_size_mb, task, return_timestamps)
        
    except Exception as e:
        logger.error(f"Error reading file: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=f"Error reading file: {str(e)}")

@app.post("/transcribe_audio_base64")
@timeit
async def transcribe_audio_base64(
    request: TranscribeAudioRequest,
    api_key: str = Depends(check_api_key)
) -> Dict[str, Any]:
    logger.debug("Starting transcribe_audio_base64 function")
    logger.debug(f"Received parameters - task: {request.task}, return_timestamps: {request.return_timestamps}")
    
    try:
        audio_data = base64.b64decode(request.audio_base64)
        file_size = len(audio_data)
        file_size_mb = file_size / (1024 * 1024)
        logger.debug(f"Decoded audio data size: {file_size} bytes ({file_size_mb:.2f}MB)")
    except Exception as e:
        logger.error(f"Error decoding base64 audio data: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=f"Error decoding base64 audio data: {str(e)}")

    return await process_audio(audio_data, file_size_mb, request.task, request.return_timestamps)

async def process_audio(audio_data: bytes, file_size_mb: float, task: str, return_timestamps: bool) -> Dict[str, Any]:
    if file_size_mb > FILE_LIMIT_MB:
        logger.warning(f"Max file size exceeded: {file_size_mb:.2f}MB > {FILE_LIMIT_MB}MB")
        raise HTTPException(status_code=400, detail=f"File size exceeds file size limit. Got file of size {file_size_mb:.2f}MB for a limit of {FILE_LIMIT_MB}MB.")

    try:
        logger.debug("Performing ffmpeg read on audio data")
        inputs = ffmpeg_read(audio_data, pipeline.feature_extractor.sampling_rate)
        inputs = {"array": inputs, "sampling_rate": pipeline.feature_extractor.sampling_rate}
        logger.debug("ffmpeg read completed successfully")
    except Exception as e:
        logger.error(f"Error in ffmpeg read: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error processing audio data: {str(e)}")

    logger.debug("Calling tqdm_generate to transcribe audio")
    try:
        text, runtime, timing_info = tqdm_generate(inputs, task=task, return_timestamps=return_timestamps)
        logger.debug(f"Transcription completed. Runtime: {runtime:.2f}s")
    except Exception as e:
        logger.error(f"Error in tqdm_generate: {str(e)}", exc_info=True)
        raise HTTPException(status_code=500, detail=f"Error transcribing audio: {str(e)}")

    logger.debug("Audio processing completed successfully")
    return jsonable_encoder({
        "text": text,
        "runtime": runtime,
        "timing_info": timing_info
    })

def tqdm_generate(inputs: dict, task: str, return_timestamps: bool):
    start_time = time.time()
    logger.debug(f"Starting tqdm_generate - task: {task}, return_timestamps: {return_timestamps}")
    
    inputs_len = inputs["array"].shape[0]
    logger.debug(f"Input array length: {inputs_len}")
    
    all_chunk_start_idx = np.arange(0, inputs_len, step)
    num_samples = len(all_chunk_start_idx)
    num_batches = math.ceil(num_samples / BATCH_SIZE)
    logger.debug(f"Number of samples: {num_samples}, Number of batches: {num_batches}")

    logger.debug("Preprocessing audio for inference")
    try:
        dataloader = pipeline.preprocess_batch(inputs, chunk_length_s=CHUNK_LENGTH_S, batch_size=BATCH_SIZE)
        logger.debug("Preprocessing completed successfully")
    except Exception as e:
        logger.error(f"Error in preprocessing: {str(e)}", exc_info=True)
        raise

    model_outputs = []
    transcription_start_time = time.time()
    logger.debug("Starting transcription...")
    
    try:
        for i, batch in enumerate(dataloader):
            logger.debug(f"Processing batch {i+1}/{num_batches} with {len(batch)} samples")
            batch_output = pipeline.forward(batch, batch_size=BATCH_SIZE, task=task, return_timestamps=True)
            model_outputs.append(batch_output)
            logger.debug(f"Batch {i+1} processed successfully")
    except Exception as e:
        logger.error(f"Error during batch processing: {str(e)}", exc_info=True)
        raise

    transcription_runtime = time.time() - transcription_start_time
    logger.debug(f"Transcription completed in {transcription_runtime:.2f}s")

    logger.debug("Post-processing transcription results")
    try:
        post_processed = pipeline.postprocess(model_outputs, return_timestamps=True)
        logger.debug("Post-processing completed successfully")
    except Exception as e:
        logger.error(f"Error in post-processing: {str(e)}", exc_info=True)
        raise

    text = post_processed["text"]
    if return_timestamps:
        timestamps = post_processed.get("chunks")
        timestamps = [
            f"[{format_timestamp(chunk['timestamp'][0])} -> {format_timestamp(chunk['timestamp'][1])}] {chunk['text']}"
            for chunk in timestamps
        ]
        text = "\n".join(str(feature) for feature in timestamps)
    
    total_processing_time = time.time() - start_time
    logger.debug("tqdm_generate function completed successfully")
    return text, transcription_runtime, {
        "transcription_time": transcription_runtime,
        "total_processing_time": total_processing_time
    }

def format_timestamp(seconds: float, always_include_hours: bool = False, decimal_marker: str = "."):
    if seconds is not None:
        milliseconds = round(seconds * 1000.0)

        hours = milliseconds // 3_600_000
        milliseconds -= hours * 3_600_000

        minutes = milliseconds // 60_000
        milliseconds -= minutes * 60_000

        seconds = milliseconds // 1_000
        milliseconds -= seconds * 1_000

        hours_marker = f"{hours:02d}:" if always_include_hours or hours > 0 else ""
        return f"{hours_marker}{minutes:02d}:{seconds:02d}{decimal_marker}{milliseconds:03d}"
    else:
        # we have a malformed timestamp so just return it as is
        return seconds