import os
import csv
import whisper
import gradio as gr
from transformers import pipeline
from fastapi import FastAPI, UploadFile, File
from fastapi.middleware.wsgi import WSGIMiddleware

# ================== DIRECTORIES =====================
UPLOAD_DIR = "./uploads"
CSV_DIR = "./csv_output"
os.makedirs(UPLOAD_DIR, exist_ok=True)
os.makedirs(CSV_DIR, exist_ok=True)

# ================== LOAD MODELS =====================
# Load Whisper (Speech-to-Text)
whisper_model = whisper.load_model("base")

# Load BART Summarizer
summarizer = pipeline(
    "summarization",
    model="facebook/bart-large-cnn"
)

# ================== FUNCTION =======================
def transcribe_and_summarize(audio_path):
    if audio_path is None:
        return "No file uploaded.", "No summary.", "API not available"

    # Transcription
    try:
        result = whisper_model.transcribe(audio_path)
        transcription = result["text"]
    except Exception as e:
        return f"Transcription failed: {e}", "No summary.", "API not available"

    # Summarization
    try:
        summary = summarizer(
            transcription, max_length=60, min_length=10, do_sample=False
        )[0]["summary_text"]
    except Exception as e:
        summary = f"Summarization failed: {e}"

    # Save to CSV
    csv_file = os.path.join(CSV_DIR, "transcription_summary.csv")
    file_exists = os.path.isfile(csv_file)
    with open(csv_file, mode="a", newline="", encoding="utf-8") as f:
        writer = csv.writer(f)
        if not file_exists:
            writer.writerow(["Audio File", "Transcription", "Summary"])
        writer.writerow([os.path.basename(audio_path), transcription, summary])

    # API URL (replace with your deployed Hugging Face Space URL)
    api_link = "https://your-username-your-space-name.hf.space/run/predict"

    return transcription, summary, api_link

# ================== GRADIO INTERFACE =================
gradio_interface = gr.Interface(
    fn=transcribe_and_summarize,
    inputs=gr.Audio(type="filepath", label="Upload Audio File"),
    outputs=[
        gr.Textbox(label="🔊 Transcription (Whisper Output)"),
        gr.Textbox(label="📝 Summary (BART Output)"),
        gr.Textbox(label="🔗 API Endpoint for Integration")
    ],
    title="🎤 Audio Transcriber + Summarizer",
    description="Upload an audio file → Transcribe it with Whisper → Summarize it using BART → Copy the API link for integration.",
)

# ================== FASTAPI APP ======================
api_app = FastAPI(title="Audio Transcriber + Summarizer API")

@api_app.post("/api/transcribe/")
async def transcribe_api(file: UploadFile = File(...)):
    temp_path = os.path.join(UPLOAD_DIR, file.filename)
    with open(temp_path, "wb") as f:
        f.write(await file.read())

    transcription, summary, api_link = transcribe_and_summarize(temp_path)
    return {"transcription": transcription, "summary": summary, "api_url": api_link}

# ================== MOUNT GRADIO ON FASTAPI ==========
api_app.mount("/gradio", WSGIMiddleware(gradio_interface.launch(prevent_thread_lock=True)))

# ================== ENTRY POINT =====================
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(api_app, host="0.0.0.0", port=8000)