import logging
import math
import time
import base64
import io
from typing import Dict, Any
from functools import wraps

from fastapi import FastAPI, Depends, HTTPException
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

# do a pre-compile step so that the first user to use the demo isn't hit with a long transcription time
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

@app.post("/transcribe_audio")
@timeit
async def transcribe_chunked_audio(
    request: TranscribeAudioRequest
) -> Dict[str, Any]:
    logger.debug("Starting transcribe_chunked_audio function")
    logger.debug(f"Received parameters - task: {request.task}, return_timestamps: {request.return_timestamps}")
    
    try:
        # Decode base64 audio data
        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)}")

    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=request.task, return_timestamps=request.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("Transcribe_chunked_audio function 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