import gradio as gr
import assemblyai as aai
from transformers import pipeline
import os
from supabase import create_client, Client
from datetime import datetime
import csv

# Add your AssemblyAI API key as Environment Variable
aai.settings.api_key = os.environ['Assembly']
url: str = os.environ['DBUrl']
key: str = os.environ['DBKey']

# Initialize question answering pipeline
question_answerer = pipeline("question-answering", model='distilbert-base-cased-distilled-squad')

# List of questions
questions = [
    "Which grade is the patient studying?",
    "How old is the patient?",
    "What is the gender?",
    "What is the chief complaint regarding the patient's oral health? If there is none, just say the word 'none' else elaborate",
    "Can you provide any relevant Medical history for the patient? If there is none, just say the word 'none', else elaborate",
    "Can you provide any relevant Dental history for the patient? If there is none, just say the word 'none', else elaborate",
    "Give me about the clinical findings listed",
    "Are there any other findings you would like to note?",
    "What treatment plan do you recommend? Choose only from Options: (Scaling, Filling, Pulp therapy/RCT, Extraction, Medication, Referral) and nothing else"
]

oral_health_assessment_form = [
    "Doctor’s Name",
    "Child’s Name",
    "Grade",
    "Age",
    "Gender",
    "Chief complaint",
    "Medical history",
    "Dental history",
    "Clinical Findings",
    "Any other finding",
    "Treatment plan",
]

# Function to generate answers for the questions
def generate_answer(question, context):
    result = question_answerer(question=question, context=context)
    return result['answer']

# Function to handle audio recording and transcription
def transcribe_audio(audio_path):
    print(f"Received audio file at: {audio_path}")
    
    # Check if the file exists and is not empty
    if not os.path.exists(audio_path):
        return "Error: Audio file does not exist."
    
    if os.path.getsize(audio_path) == 0:
        return "Error: Audio file is empty."
    
    try:
        # Transcribe the audio file using AssemblyAI
        transcriber = aai.Transcriber()
        print("Starting transcription...")
        transcript = transcriber.transcribe(audio_path)
        
        print("Transcription process completed.")
        
        # Handle the transcription result
        if transcript.status == aai.TranscriptStatus.error:
            print(f"Error during transcription: {transcript.error}")
            return transcript.error
        else:
            context = transcript.text
            print(f"Transcription text: {context}")
            return context
    
    except Exception as e:
        print(f"Exception occurred: {e}")
        return str(e)

# Function to fill in the answers for the text boxes
def fill_textboxes(context):
    answers = []
    for question in questions:
        answer = generate_answer(question, context)
        answers.append(answer)
    return answers

# Supabase configuration
supabase: Client = create_client(url, key)

# Main Gradio app function
def main(audio):
    context = transcribe_audio(audio)
    
    if "Error" in context:
        return [context] * len(oral_health_assessment_form)
    
    answers = fill_textboxes(context)
    answers.insert(0, "")  # for "Child's Name"
    answers.insert(0, "Dr.Alvi")  # for "Doctor’s Name"
    
    return answers

def save_answers(*args):
    current_datetime = datetime.now().isoformat()
    answers_dict = {oral_health_assessment_form[i]: args[i] for i in range(len(oral_health_assessment_form))}
    answers_dict["Submission Date and Time"] = current_datetime
    print("Saved answers:", answers_dict)
    
    # Insert data into Supabase
    try:
        response = supabase.table('oral_health_assessments').insert(answers_dict).execute()
        print("Data inserted into Supabase:", response.data)
        return f"Saved answers: {answers_dict}"
    except Exception as e:
        print(f"Error inserting data into Supabase: {e}")
        return f"Error saving answers: {e}"

# Function to download table as CSV
def download_table_to_csv():
    # Fetch data from Supabase table
    response = supabase.table("oral_health_assessments").select("*").execute()
    
    # Check if data is available
    if not response.data:
        print("No data found in the table.")
        return None

    data = response.data

    # Prepare CSV data
    csv_data = []

    # Add header row
    if len(data) > 0:
        csv_data.append(data[0].keys())

    # Add data rows
    for row in data:
        csv_data.append(row.values())

    # Save CSV data to file (replace 'your_table.csv' with desired filename)
    csv_file = "your_table.csv"
    with open(csv_file, "w", newline='') as f:
        writer = csv.writer(f)
        writer.writerows(csv_data)

    print("Downloaded table oral_health_assessments")
    return csv_file

def gradio_download():
    file_path = download_table_to_csv()
    if file_path:
        return file_path
    return None

# Create the Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# OHA Form Filler App")
    audio_input = gr.Audio(type="filepath", label="Record your audio", elem_id="audio_input")
    transcribe_button = gr.Button("Transcribe and Generate Form", elem_id="transcribe_button")
    
    with gr.Row(elem_id="textboxes_row"):
        with gr.Column():
            textboxes_left = [gr.Textbox(label=oral_health_assessment_form[i], value="", interactive=True) for i in range(0, len(oral_health_assessment_form)//2)]
        with gr.Column():
            textboxes_right = [gr.Textbox(label=oral_health_assessment_form[i], value="", interactive=True) for i in range(len(oral_health_assessment_form)//2, len(oral_health_assessment_form))]
    
    submit_button = gr.Button("Submit", elem_id="submit_button")
    output_html = gr.HTML(label="Submitted Answers")

    transcribe_button.click(fn=main, inputs=audio_input, outputs=textboxes_left + textboxes_right)
    submit_button.click(fn=save_answers, inputs=textboxes_left + textboxes_right, outputs=output_html)

    download_button = gr.Button("Download Table as CSV")
    download_csv_output = gr.File(label="Download CSV")

    download_button.click(fn=gradio_download, inputs=[], outputs=download_csv_output)

# Launch the app
demo.launch(share=True)