import torch
import whisper
from transformers import pipeline
import gradio as gr
import concurrent.futures
import os  # For environment variables

print("Starting up...")

# *** Model Loading - CPU Optimized & Size Considerations ***
try:
    # Option 1: Try "tiny" model.  Significantly faster on CPU, but lower accuracy.
    whisper_model = whisper.load_model("tiny")
    print("Using whisper 'tiny' model.")
except Exception as e:
    print(f"Error loading whisper 'tiny' model: {e}.  Trying 'small'.")
    try:
        whisper_model = whisper.load_model("small")
        print("Using whisper 'small' model.")
    except Exception as e2:
        print(f"Error loading whisper 'small' model: {e2}.  Whisper will not work.")
        whisper_model = None  # Disable whisper functionality

try:
    summarizer = pipeline("summarization", model="facebook/bart-large-cnn", device=-1)  # device=-1 forces CPU
    question_generator = pipeline("text2text-generation", model="google/flan-t5-large", device=-1) # device=-1 forces CPU
    print("Summarizer and Question Generator loaded successfully.")

except Exception as e:
    print(f"Error loading Summarizer or Question Generator: {e}")
    summarizer = None
    question_generator = None
    print("Summarization and Question Generation will not work.")


print("Models loaded (or failed gracefully).")


# *** Transcription ***
def transcribe_audio(audio_path):
    print("Transcribing audio...")
    if whisper_model is None:
        return "Error: Whisper model failed to load."

    try:
        result = whisper_model.transcribe(audio_path)
        return result["text"]
    except Exception as e:
        print(f"Error transcribing audio: {e}")
        return f"Error during transcription: {e}"


# *** Summarization ***
def summarize_text(text):
    if summarizer is None:
        return "Error: Summarizer model failed to load."

    print("Summarizing text using BART...")
    # Chunk the text into smaller parts, even smaller than before for CPU
    text_chunks = [text[i:i + 768] for i in range(0, len(text), 768)]  # More aggressive chunking

    try:
        summaries = summarizer(text_chunks, max_length=150, min_length=30, do_sample=False)  # Reduce length
        return " ".join([s['summary_text'] for s in summaries])
    except Exception as e:
        print(f"Error during summarization: {e}")
        return f"Error during summarization: {e}"


# *** Question Generation ***
def generate_questions(text):
    if question_generator is None:
        return "Error: Question Generation model failed to load."

    print("Generating questions using FLAN-T5...")
    # Even smaller chunks for question generation (CPU is struggling)
    text_chunks = [text[i:i + 512] for i in range(0, len(text), 512)]

    questions = []

    with concurrent.futures.ThreadPoolExecutor(max_workers=os.cpu_count()) as executor: # Explicitly limit threads
        future_questions = [
            executor.submit(
                lambda chunk: question_generator(
                    f"You are an AI tutor. Your task is to generate **insightful, topic-specific** questions based on the following passage. Ensure that the questions are relevant to the **key concepts, definitions, and explanations** present in the text. Avoid generic questions.\n\nPassage:\n{chunk}",
                    max_length=80, num_return_sequences=2, do_sample=True  # Reduce length and sequences
                ),
                chunk
            ) for chunk in text_chunks
        ]

        for future in future_questions:
            try:
                generated = future.result()
                questions.extend([q['generated_text'] for q in generated])
            except Exception as e:
                print(f"Error generating questions for a chunk: {e}")
                questions.append(f"Error generating questions: {e}")  # Report the error

    return "\n".join(questions)


# *** Main Processing Function ***
def process_audio(audio_path):
    transcript = transcribe_audio(audio_path)
    summary = summarize_text(transcript)
    questions = generate_questions(transcript)

    combined_text = f"📝 Transcription:\n{transcript}\n\n📜 Summary:\n{summary}\n\n🤔 Practice Questions:\n{questions}"
    file_path = "lecture_summary.txt"

    with open(file_path, "w", encoding="utf-8") as f:
        f.write(combined_text)

    return transcript, summary, questions, file_path

def gradio_interface(audio):
    return process_audio(audio)



# *** Gradio Interface ***
with gr.Blocks(css="""
    #submit-btn, #download-btn {
        background-color: blue !important;
        color: white !important;
        border-radius: 8px !important;
        padding: 10px !important;
        font-size: 16px !important;
    }
    textarea {
        border: 2px solid black !important;
        border-radius: 5px !important;
    }
""") as demo:
    gr.Markdown("# 🎙 LectureGenie: Transcribe, Summarize & Quiz")
    gr.Markdown("Upload a lecture audio file. The system will **transcribe**, **summarize**, and **generate questions** automatically.")

    audio_input = gr.Audio(type="filepath", label="🎤 Upload Audio File", interactive=True)
    submit_btn = gr.Button("Submit", elem_id="submit-btn")

    with gr.Row():
        with gr.Column():
            transcript_box = gr.Textbox(label="📝 Transcription", lines=10, interactive=False, show_copy_button=True)

        with gr.Column():
            summary_box = gr.Textbox(label="📜 Summary", lines=10, interactive=False, show_copy_button=True)

        with gr.Column():
            questions_box = gr.Textbox(label="🤔 Practice Questions", lines=10, interactive=False, show_copy_button=True)

    download_btn = gr.File(label="📥 Download All", interactive=False, visible=False)
    download_button = gr.Button("📥 Download", elem_id="download-btn")

    submit_btn.click(
        gradio_interface,
        inputs=[audio_input],
        outputs=[transcript_box, summary_box, questions_box, download_btn]
    )

    download_button.click(lambda x: x, inputs=[download_btn], outputs=[download_btn])

demo.launch(share=True)