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

from fastapi import FastAPI, Depends, HTTPException, File, UploadFile
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 = None
    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():
    api_key = os.environ.get("WHISPER_API_KEY")
    if not api_key:
        raise HTTPException(status_code=401, detail="API key not found in environment variables")
    return api_key

@app.post("/transcribe_audio")
@timeit
async def transcribe_chunked_audio(
    request: TranscribeAudioRequest = None,
    file: UploadFile = File(None),
    api_key: str = Depends(check_api_key)
) -> Dict[str, Any]:
    logger.debug("Starting transcribe_chunked_audio function")
    logger.debug(f"Received parameters - task: {request.task if request else 'transcribe'}, return_timestamps: {request.return_timestamps if request else False}")
    
    try:
        if file:
            logger.debug("Processing uploaded file")
            audio_data = await file.read()
            file_size = len(audio_data)
        elif request and request.audio_base64:
            logger.debug("Processing base64 encoded audio")
            audio_data = base64.b64decode(request.audio_base64)
            file_size = len(audio_data)
        else:
            raise HTTPException(status_code=400, detail="No audio data provided")
        
        file_size_mb = file_size / (1024 * 1024)
        logger.debug(f"Audio data size: {file_size} bytes ({file_size_mb:.2f}MB)")
    except Exception as e:
        logger.error(f"Error processing audio data: {str(e)}", exc_info=True)
        raise HTTPException(status_code=400, detail=f"Error processing 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:
        task = request.task if request else "transcribe"
        return_timestamps = request.return_timestamps if request else False
        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("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