from fastapi import FastAPI, Request, Response, status, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel
import torch
import io
import base64
import soundfile as sf
from transformers import VitsModel, AutoTokenizer
import logging
import torch.quantization # Import the quantization library
import ffmpeg

# Configure basic logging to console
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)
print("I AM IN BACKEND")
app = FastAPI()

# --- Custom Logging Middleware for OPTIONS Requests (for debugging 400 Bad Request) ---
@app.middleware("http")
async def log_options_requests(request: Request, call_next):
    """
    Logs details of OPTIONS requests to help diagnose 400 Bad Request errors.
    This middleware runs before CORSMiddleware.
    """
    if request.method == "OPTIONS":
        logger.info(f"--- Received OPTIONS Request for Path: {request.url.path} ---")
        logger.info(f"Client Host: {request.client.host}:{request.client.port}")
        for header, value in request.headers.items():
            logger.info(f"Header: {header}: {value}")
        logger.info("-----------------------------------------------------")
    
    # Process the request with the next middleware/route handler
    response = await call_next(request)
    return response
# --- End Custom Logging Middleware ---


# --- CORS Configuration ---
# Define the origins that are allowed to make cross-origin requests.
# It's crucial for the frontend (Flutter web app) to be listed here.
# In production, replace specific ports/local IPs with your deployed frontend URL.
import re # Import regex for more flexible origin matching
# ...
origins = [
    "http://localhost",
    "http://localhost:42389", # Common Flutter web release port
    "http://127.0.0.1",
    "http://127.0.0.1:9101",  # Flutter DevTools origin
    # Regex patterns to allow any port on localhost and 127.0.0.1 for development
    re.compile(r"^http:\/\/localhost:\d+$"),
    re.compile(r"^http:\/\/127\.0\.0\.1:\d+$"),
    # ...
]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,        # List of allowed origins
    allow_credentials=True,       # Allow credentials (e.g., cookies, authorization headers)
    allow_methods=["*"],          # Allow all HTTP methods (GET, POST, PUT, DELETE, OPTIONS, etc.)
    allow_headers=["*"],          # Allow all headers in the request
)
# --- End CORS Configuration ---


# --- Text-to-Speech Model Loading ---
# --- Text-to-Speech Model Loading ---

# Load model once when the application starts
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model_name = "Benjamin-png/swahili-mms-tts-finetuned"
try:
    model = VitsModel.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    
    # Check if a GPU is not available, and if so, apply quantization.
    if device.type == 'cpu':
        logger.info("GPU not found. Applying dynamic quantization for CPU optimization.")
        # Eager mode quantization
        model = torch.quantization.quantize_dynamic(
            model, 
            {torch.nn.Linear}, 
            dtype=torch.qint8
        )
    
    # Now move the potentially quantized model to the device
    model.to(device)
    logger.info("TTS model and tokenizer loaded successfully.")

except Exception as e:
    logger.error(f"Failed to load TTS model or tokenizer: {e}")
# --- End TTS Model Loading ---
# Load model once when the application starts
# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# model_name = "Benjamin-png/swahili-mms-tts-finetuned"
# try:
#     model = VitsModel.from_pretrained(model_name).to(device)
#     tokenizer = AutoTokenizer.from_pretrained(model_name)
#     logger.info("TTS model and tokenizer loaded successfully.")
# except Exception as e:
#     logger.error(f"Failed to load TTS model or tokenizer: {e}")
#     # Consider raising an exception or having a fallback here if model loading is critical
    # For now, the app will start but TTS calls will fail.
# --- End TTS Model Loading ---


# --- Request Body Model ---
# Defines the expected structure for the incoming JSON request
class TTSRequest(BaseModel):
    text: str
# --- End Request Body Model ---


# --- FastAPI Endpoint for TTS ---
@app.post("/tts/")
async def generate_speech(request: TTSRequest):
    """
    Generates speech from the provided text using a pre-trained VITS model.
    The audio is returned as a base64 encoded WAV file.
    """
    if not model or not tokenizer:
        raise HTTPException(status_code=500, detail="TTS model not loaded.")

    try:
        # Tokenize the input text and move it to the appropriate device (CPU/GPU)
        inputs = tokenizer(request.text, return_tensors="pt").to(device)

        # Perform inference without tracking gradients for efficiency
        with torch.no_grad():
            output = model(**inputs).waveform

        # Move the output audio data to CPU and convert to a NumPy array
        audio = output.squeeze().cpu().numpy()

        # Create an in-memory buffer to store the WAV file
        # buffer = io.BytesIO()
        # # Write the audio data to the buffer in WAV format
        # sf.write(buffer, audio, samplerate=model.config.sampling_rate, format="WAV")
        # # Reset the buffer's position to the beginning to read its content
        # buffer.seek(0)

        # # Read the buffer content, base64 encode it, and decode to a UTF-8 string
        # audio_base64 = base64.b64encode(buffer.read()).decode('utf-8')

        # # Return the base64 encoded audio in a JSON response
        # return {"audio_base64": audio_base64}
        # Create an in-memory buffer to store the intermediate WAV file
        buffer_wav = io.BytesIO()
        sf.write(buffer_wav, audio, samplerate=model.config.sampling_rate, format="WAV")
        buffer_wav.seek(0)

        # Use ffmpeg to convert the WAV data to MP3 data in memory
        try:
            # The 'run' method returns the standard output, which is our MP3 data
            audio_mp3_bytes, _ = (
                ffmpeg
                .input('pipe:0', format='wav')
                .output('pipe:1', format='mp3', audio_bitrate='64k') # Adjust bitrate as needed
                .run(input=buffer_wav.read(), capture_stdout=True, capture_stderr=True)
            )
        except ffmpeg.Error as e:
            logger.error(f"FFmpeg error: {e.stderr.decode()}")
            raise HTTPException(status_code=500, detail="Audio conversion failed.")

        # Base64 encode the MP3 bytes
        audio_base64 = base64.b64encode(audio_mp3_bytes).decode('utf-8')

        # Return the base64 encoded audio in a JSON response
        return {"audio_base64": audio_base64}
    except Exception as e:
        logger.error(f"Error during speech generation: {e}")
        raise HTTPException(status_code=500, detail=f"Speech generation failed: {e}")
# --- End TTS Endpoint ---


# Example root endpoint (optional)
@app.get("/")
def read_root():
    """
    A simple root endpoint to confirm the backend is running.
    """
    return {"Hello": "From FastAPI Backend"}