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 child studying?",
    "How old is the child?",
    "What is the gender?",
    "Can you provide the name and location of the child's school?",
    "What are the names of the child's guardians or parents?",
    "What is the chief complaint regarding the child's oral health? If there is none, just say the word 'none', else elaborate only on medication history",
    "Can you provide any relevant medical history for the child? If there is none, just say the word 'none', else elaborate",
    "Does the child take any medications regularly? If there is none, just say the word 'none'. If yes, please specify.",
    "When was the child's previous dental visit? If no visits before, just say the word 'first' or mention the visit number and nothing else",
    "Does the child have any habits such as thumb sucking, tongue thrusting, nail biting, or lip biting? If yes, just list them and don't provide any further details",
    "Does the patient brush their teeth? Just use the words 'once daily', 'twice daily', or 'thrice daily' to answer, nothing else",
    "Does the child experience bleeding gums? Just say 'yes' or 'no' for this and nothing else",
    "Has the child experienced early childhood caries? Just say 'yes' or 'no' and nothing else",
    "Please mention if tooth decay is present with tooth number(s), else just say the word 'none' and nothing else",
    "Have any teeth been fractured? If yes, please mention the tooth number(s), else just say 'none' and nothing else",
    "Is there any pre-shedding mobility of teeth? If yes, please specify, else just say 'none' and nothing else",
    "Does the child have malocclusion? If yes, please provide details, else just say the word 'none' and nothing else",
    "Does the child experience pain, swelling, or abscess? If yes, please provide details, else just say 'none' and nothing else",
    "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",
    "School name and place",
    "Guardian/Parents name",
    "Chief complaint",
    "Medical history",
    "Medication history",
    "Previous dental visit",
    "Habits",
    "Brushing habit",
    "Bleeding gums",
    "Early Childhood caries",
    "Decayed",
    "Fractured teeth",
    "Preshedding mobility",
    "Malocclusion",
    "Does the child have pain, swelling or abscess?",
    "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 "No data found in the table."
    
    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
    file_path = "oral_health_assessments.csv"
    with open(file_path, "w", newline='') as f:
        writer = csv.writer(f)
        writer.writerows(csv_data)

    print(f"Downloaded table 'oral_health_assessments' to '{file_path}'")
    return file_path

# 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")
    download_button = gr.Button("Download as CSV", elem_id="download_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.click(fn=download_table_to_csv, outputs=gr.File())

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