import gradio as gr
import torch
from transformers import BertTokenizerFast, BertForQuestionAnswering, pipeline
import torch.nn.functional as F
import speech_recognition as sr

# Load QA Model
tokenizer = BertTokenizerFast.from_pretrained("bert-base-uncased")
model = BertForQuestionAnswering.from_pretrained("bert-base-uncased")
model.load_state_dict(torch.load("lexpilot_trained_model.pt", map_location=torch.device("cpu")))
model.eval()

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

def answer_with_confidence(question, context):
    inputs = tokenizer.encode_plus(question, context, return_tensors="pt")
    with torch.no_grad():
        outputs = model(**inputs)
        start_logits = outputs.start_logits
        end_logits = outputs.end_logits
        start_probs = F.softmax(start_logits, dim=1)
        end_probs = F.softmax(end_logits, dim=1)
        start = torch.argmax(start_probs)
        end = torch.argmax(end_probs) + 1
        confidence = (start_probs[0][start] * end_probs[0][end - 1]).item()
        tokens = inputs["input_ids"][0][start:end]
        answer = tokenizer.decode(tokens)
    return answer.strip(), round(confidence, 2)

def summarize_text(text):
    summary = summarizer(text, max_length=100, min_length=30, do_sample=False)
    return summary[0]["summary_text"]

def voice_to_text(audio):
    recognizer = sr.Recognizer()
    with sr.AudioFile(audio) as source:
        audio_data = recognizer.record(source)
    try:
        return recognizer.recognize_google(audio_data)
    except Exception as e:
        return f"Error: {str(e)}"

def qa_interface(context, question):
    answer, confidence = answer_with_confidence(question, context)
    summary = summarize_text(context)
    return answer, confidence, summary

with gr.Blocks() as app:
    gr.Markdown("## 🤖 LexPilot: QA + Summarization + Voice Input")

    with gr.Row():
        context_input = gr.Textbox(lines=6, label="📘 Context")
        question_input = gr.Textbox(lines=1, label="❓ Question")

    with gr.Row():
        answer_output = gr.Textbox(label="🧠 Answer")
        confidence_output = gr.Number(label="📊 Confidence")
        summary_output = gr.Textbox(label="📝 Summary")

    with gr.Row():
        submit_btn = gr.Button("🔍 Get Answer & Summary")
        audio_input = gr.Audio(type="filepath", label="🎙️ Upload Voice Question")

    submit_btn.click(fn=qa_interface, inputs=[context_input, question_input],
                     outputs=[answer_output, confidence_output, summary_output])

    audio_input.change(fn=voice_to_text, inputs=audio_input, outputs=question_input)

app.launch()