# Complete Malayalam Hospital Booking Chatbot using Llama 3.1-8B-Instruct
# with HuggingFace Transformers Library in Google Colab


import json
import gradio as gr
import torch
import datetime
import pytz
import uuid
import re
import time
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
from google.colab import auth
from googleapiclient.discovery import build
import os

# Set up timezone for India
IST = pytz.timezone('Asia/Kolkata')

# ===== CONFIGURATION =====

# Path to store the model locally (to avoid re-downloading)
MODEL_PATH = "/content/llama-3.1-8b-instruct"

# Replace with your actual Hugging Face token
HF_TOKEN = ""  # Will be set via Colab input

# Google Calendar API scopes
SCOPES = ['https://www.googleapis.com/auth/calendar']

# Available doctors and departments for booking
available_doctors = {
  "cardiology": ["Dr. Anoop Menon", "Dr. Priya Nair"],
  "orthopedics": ["Dr. Rajesh Kumar", "Dr. Meera Pillai"],
  "neurology": ["Dr. Vinod Thomas", "Dr. Lakshmi Nair"],
  "pediatrics": ["Dr. Suresh Babu", "Dr. Anjali Krishnan"],
  "general": ["Dr. Joseph Mathew", "Dr. Deepa Varma"]
}

# Hospital database simulation
appointments_db = {}

# ===== FUNCTION DEFINITIONS =====

function_definitions = [
  {
    "name": "check_doctor_availability",
    "description": "Check which doctors are available in a specific department",
    "parameters": {
      "type": "object",
      "properties": {
        "department": {
          "type": "string",
          "description": "The hospital department (cardiology, orthopedics, neurology, pediatrics, general)"
        }
      },
      "required": ["department"]
    }
  },
  {
    "name": "check_time_slots",
    "description": "Check available time slots for a specific doctor on a specific date",
    "parameters": {
      "type": "object",
      "properties": {
        "doctor_name": {
          "type": "string",
          "description": "The name of the doctor"
        },
        "date": {
          "type": "string",
          "description": "The date in YYYY-MM-DD format"
        }
      },
      "required": ["doctor_name", "date"]
    }
  },
  {
    "name": "book_appointment",
    "description": "Book an appointment with a doctor and add it to Google Calendar",
    "parameters": {
      "type": "object",
      "properties": {
        "patient_name": {
          "type": "string",
          "description": "The name of the patient"
        },
        "patient_phone": {
          "type": "string",
          "description": "The phone number of the patient"
        },
        "doctor_name": {
          "type": "string",
          "description": "The name of the doctor"
        },
        "department": {
          "type": "string",
          "description": "The hospital department"
        },
        "date": {
          "type": "string",
          "description": "The date in YYYY-MM-DD format"
        },
        "time": {
          "type": "string",
          "description": "The time of the appointment (e.g., '10:00 AM')"
        },
        "description": {
          "type": "string",
          "description": "Brief description of the medical issue"
        }
      },
      "required": ["patient_name", "patient_phone", "doctor_name", "department", "date", "time"]
    }
  },
  {
    "name": "cancel_appointment",
    "description": "Cancel an existing appointment",
    "parameters": {
      "type": "object",
      "properties": {
        "appointment_id": {
          "type": "string",
          "description": "The ID of the appointment to cancel"
        },
        "patient_phone": {
          "type": "string",
          "description": "The phone number of the patient for verification"
        }
      },
      "required": ["appointment_id", "patient_phone"]
    }
  }
]

# ===== FUNCTION IMPLEMENTATIONS =====

def check_doctor_availability(department):
    """Check which doctors are available in a specific department"""
    if department.lower() in available_doctors:
        return {
            "available": True,
            "doctors": available_doctors[department.lower()]
        }
    else:
        return {
            "available": False,
            "message": "Department not found",
            "available_departments": list(available_doctors.keys())
        }

def check_time_slots(doctor_name, date):
    """Check available time slots for a specific doctor on a specific date"""
    # Available time slots
    all_slots = [
        "09:00 AM", "09:30 AM", "10:00 AM", "10:30 AM",
        "11:00 AM", "11:30 AM", "12:00 PM", "02:00 PM",
        "02:30 PM", "03:00 PM", "03:30 PM", "04:00 PM"
    ]

    # In a real implementation, this would check a database
    # For this example, we'll simulate some slots being taken
    taken_slots = random.sample(all_slots, 3)  # Randomly mark 3 slots as taken

    available_slots = [slot for slot in all_slots if slot not in taken_slots]

    return {
        "date": date,
        "doctor": doctor_name,
        "available_slots": available_slots
    }

def book_appointment(appointment_details, calendar_service):
    """Book an appointment with a doctor and add it to Google Calendar"""
    try:
        # Validate the appointment details first
        doctor_exists = False
        for dept_doctors in available_doctors.values():
            if appointment_details["doctor_name"] in dept_doctors:
                doctor_exists = True
                break

        if not doctor_exists:
            return {
                "success": False,
                "message": "Doctor not found"
            }

        # Parse date and time
        date_parts = appointment_details["date"].split('-')
        year, month, day = int(date_parts[0]), int(date_parts[1]), int(date_parts[2])

        time_parts = appointment_details["time"].split(' ')
        time = time_parts[0]
        meridian = time_parts[1]

        hours, minutes = map(int, time.split(':'))

        if meridian == 'PM' and hours != 12:
            hours += 12
        if meridian == 'AM' and hours == 12:
            hours = 0

        start_datetime = datetime.datetime(year, month, day, hours, minutes, 0, tzinfo=IST)
        end_datetime = start_datetime + datetime.timedelta(minutes=30)  # 30 minutes appointment

        # Create the calendar event
        event = {
            'summary': f"Medical appointment with {appointment_details['doctor_name']}",
            'location': 'City Hospital, Kochi, Kerala',
            'description': appointment_details.get('description', 'Regular checkup'),
            'start': {
                'dateTime': start_datetime.isoformat(),
                'timeZone': 'Asia/Kolkata',
            },
            'end': {
                'dateTime': end_datetime.isoformat(),
                'timeZone': 'Asia/Kolkata',
            },
            'attendees': [
                {'email': 'doctor@cityhospital.com'},
                {'email': 'patient@example.com'}  # In a real app, use actual email
            ],
            'reminders': {
                'useDefault': False,
                'overrides': [
                    {'method': 'email', 'minutes': 24 * 60},
                    {'method': 'popup', 'minutes': 60},
                ],
            },
        }

        # Add to Google Calendar
        if calendar_service:
            try:
                event = calendar_service.events().insert(calendarId='primary', body=event).execute()
                appointment_id = event['id']
            except Exception as e:
                print(f"Calendar service error: {e}")
                # Generate a mock ID if calendar service fails
                appointment_id = str(uuid.uuid4())
        else:
            # If no calendar service, generate a mock ID
            appointment_id = str(uuid.uuid4())

        # Store in local database
        appointments_db[appointment_id] = {
            "patient_name": appointment_details["patient_name"],
            "patient_phone": appointment_details["patient_phone"],
            "doctor_name": appointment_details["doctor_name"],
            "department": appointment_details["department"],
            "date": appointment_details["date"],
            "time": appointment_details["time"],
            "description": appointment_details.get("description", ""),
        }

        return {
            "success": True,
            "appointment_id": appointment_id,
            "message": "Appointment successfully booked",
            "details": {
                "doctor": appointment_details["doctor_name"],
                "department": appointment_details["department"],
                "date": appointment_details["date"],
                "time": appointment_details["time"],
                "location": 'City Hospital, Kochi, Kerala'
            }
        }
    except Exception as e:
        print(f"Error in book_appointment: {e}")
        return {
            "success": False,
            "message": f"Failed to book appointment: {str(e)}"
        }

def cancel_appointment(appointment_id, patient_phone, calendar_service):
    """Cancel an existing appointment"""
    try:
        # Check if appointment exists in our database
        if appointment_id not in appointments_db:
            return {
                "success": False,
                "message": "Appointment not found"
            }

        # Verify patient phone
        if appointments_db[appointment_id]["patient_phone"] != patient_phone:
            return {
                "success": False,
                "message": "Patient phone number does not match our records"
            }

        # Delete from Google Calendar
        if calendar_service:
            try:
                calendar_service.events().delete(calendarId='primary', eventId=appointment_id).execute()
            except Exception as e:
                print(f"Error deleting from calendar: {e}")
                # Continue anyway to delete from local database

        # Remove from local database
        del appointments_db[appointment_id]

        return {
            "success": True,
            "message": "Appointment successfully cancelled"
        }
    except Exception as e:
        return {
            "success": False,
            "message": f"Failed to cancel appointment: {str(e)}"
        }

# ===== GOOGLE CALENDAR AUTHENTICATION =====

def get_calendar_service():
    """Authenticate and return the Google Calendar service"""
    creds = None

    try:
        # Authenticate using Colab's auth helper
        auth.authenticate_user()

        # Get credentials from the authenticated Colab user
        from google.auth import default
        creds, _ = default()

        # Build and return the service
        service = build('calendar', 'v3', credentials=creds)
        return service
    except Exception as e:
        print(f"Error authenticating with Google Calendar: {e}")
        print("Continuing without Google Calendar integration.")
        return None

# ===== LLAMA 3.1 MODEL SETUP =====

def load_llama_model():
    """Load the Llama 3.1 model and tokenizer"""
    model_name = "meta-llama/Meta-Llama-3.1-8B-Instruct"

    print("Loading Llama 3.1 model and tokenizer...")

    try:
        # Check if model is already downloaded
        if os.path.exists(MODEL_PATH):
            print(f"Loading model from local path: {MODEL_PATH}")
            tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH)
            model = AutoModelForCausalLM.from_pretrained(
                MODEL_PATH,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                low_cpu_mem_usage=True
            )
        else:
            print(f"Downloading model from Hugging Face Hub")
            tokenizer = AutoTokenizer.from_pretrained(model_name, token=HF_TOKEN)
            model = AutoModelForCausalLM.from_pretrained(
                model_name,
                torch_dtype=torch.bfloat16,
                device_map="auto",
                low_cpu_mem_usage=True,
                token=HF_TOKEN
            )

            # Save model locally to avoid re-downloading
            print(f"Saving model to: {MODEL_PATH}")
            tokenizer.save_pretrained(MODEL_PATH)
            model.save_pretrained(MODEL_PATH)

        print("Model loaded successfully!")
        return model, tokenizer

    except Exception as e:
        print(f"Error loading model: {e}")
        return None, None

# ===== CHAT PROCESSING =====

def format_prompt_with_functions(messages, system_prompt):
    """Format the prompt for Llama 3.1 with function definitions"""
    # Add function definitions to system prompt
    full_system_prompt = system_prompt + "\n\n"
    full_system_prompt += "You have access to the following functions:\n"

    for func in function_definitions:
        full_system_prompt += f"- {func['name']}: {func['description']}\n"
        full_system_prompt += "  Parameters:\n"
        for param_name, param_info in func['parameters']['properties'].items():
            required = "required" if param_name in func['parameters'].get('required', []) else "optional"
            full_system_prompt += f"    - {param_name} ({required}): {param_info.get('description', '')}\n"

    full_system_prompt += "\nIf the user's request can be addressed by calling one of these functions, respond in the following JSON format:\n"
    full_system_prompt += '```json\n{"function_call": {"name": "function_name", "arguments": {"arg1": "value1", "arg2": "value2"}}}\n```\n'
    full_system_prompt += "Otherwise, respond conversationally."

    # Format conversation history
    formatted_messages = [
        {"role": "system", "content": full_system_prompt}
    ]

    # Add conversation history
    for message in messages:
        if message["role"] == "function":
            # Convert function results to assistant format for Llama 3.1
            formatted_messages.append({
                "role": "assistant",
                "content": f"I'll process the function result: {message['content']}"
            })
        else:
            formatted_messages.append(message)

    return formatted_messages

def extract_function_call(response_text):
    """Extract function call from model response"""
    # Look for JSON block in the response
    json_pattern = r'```json\s*(.*?)\s*```'
    json_matches = re.findall(json_pattern, response_text, re.DOTALL)

    if not json_matches:
        # Try alternative pattern without markdown
        json_pattern = r'({.*"function_call".*})'
        json_matches = re.findall(json_pattern, response_text, re.DOTALL)

    if json_matches:
        try:
            for json_str in json_matches:
                parsed_json = json.loads(json_str.strip())
                if "function_call" in parsed_json:
                    function_call = parsed_json["function_call"]
                    return {
                        "id": str(uuid.uuid4()),
                        "name": function_call["name"],
                        "arguments": function_call["arguments"]
                    }
        except json.JSONDecodeError:
            print(f"Failed to parse JSON: {json_matches[0]}")

    return None

def process_chat(message, chat_history, language, model_tokenizer_calendar):
    """Process a chat message, calling functions when necessary"""
    model, tokenizer, calendar_service = model_tokenizer_calendar

    try:
        # Create system prompt based on language preference
        system_prompt = f"""You are a hospital booking assistant for City Hospital in Kerala. You can understand and respond fluently in Malayalam and English.

For Malayalam speakers, introduce yourself as: "ഹലോ, ഞാൻ സിറ്റി ഹോസ്പിറ്റലിന്റെ ഓൺലൈൻ അസിസ്റ്റന്റ് ആണ്. എങ്ങനെ സഹായിക്കാൻ കഴിയും?"

Be polite and helpful. You can assist with checking doctor availability, booking appointments, and answering general questions about the hospital services.

For medical questions that require diagnosis, always advise patients to consult a doctor directly.

When booking appointments, collect all necessary information: patient name, phone number, department, doctor, date, and time.

Current language preference: {language}"""

        # Build message history from chat history
        messages = []
        for user_msg, bot_msg in chat_history:
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": bot_msg})

        # Add current message
        messages.append({"role": "user", "content": message})

        # Format messages with function calling info
        formatted_messages = format_prompt_with_functions(messages, system_prompt)

        # Generate model response
        inputs = tokenizer.apply_chat_template(
            formatted_messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt"
        ).to(model.device)

        outputs = model.generate(
            inputs,
            max_new_tokens=1024,
            temperature=0.7,
            top_p=0.9,
            do_sample=True
        )

        response_text = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)

        # Check if response contains a function call
        function_call = extract_function_call(response_text)

        if function_call:
            # Extract non-JSON response part for context (if any)
            response_context = response_text.split("```")[0].strip() if "```" in response_text else ""

            # Execute the appropriate function
            function_name = function_call["name"]
            function_args = function_call["arguments"]

            function_result = None
            if function_name == "check_doctor_availability" and "department" in function_args:
                function_result = check_doctor_availability(function_args["department"])
            elif function_name == "check_time_slots" and "doctor_name" in function_args and "date" in function_args:
                function_result = check_time_slots(function_args["doctor_name"], function_args["date"])
            elif function_name == "book_appointment":
                function_result = book_appointment(function_args, calendar_service)
            elif function_name == "cancel_appointment" and "appointment_id" in function_args and "patient_phone" in function_args:
                function_result = cancel_appointment(function_args["appointment_id"], function_args["patient_phone"], calendar_service)
            else:
                function_result = {"error": "Invalid function call or missing parameters"}

            # Add the function result to messages
            messages.append({
                "role": "assistant",
                "content": response_context,
            })

            messages.append({
                "role": "function",
                "name": function_name,
                "content": json.dumps(function_result)
            })

            # Format messages for second call
            formatted_messages = format_prompt_with_functions(messages, system_prompt)

            # Generate second response
            inputs = tokenizer.apply_chat_template(
                formatted_messages,
                tokenize=True,
                add_generation_prompt=True,
                return_tensors="pt"
            ).to(model.device)

            outputs = model.generate(
                inputs,
                max_new_tokens=1024,
                temperature=0.7,
                top_p=0.9,
                do_sample=True
            )

            second_response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)

            # Update chat history
            new_chat_history = chat_history + [(message, second_response)]

            return second_response, new_chat_history
        else:
            # No function call, just return the response
            # Update chat history
            new_chat_history = chat_history + [(message, response_text)]

            return response_text, new_chat_history

    except Exception as e:
        print(f"Error in process_chat: {e}")
        error_msg = f"Sorry, I encountered an error. Please try again. (Error: {str(e)})"
        return error_msg, chat_history + [(message, error_msg)]

# ===== GRADIO INTERFACE =====

def create_gradio_interface(model, tokenizer, calendar_service):
    """Create the Gradio interface for the chatbot"""

    with gr.Blocks(css="""
        .gradio-container {max-width: 800px !important}
        .chat-window {height: 600px !important; overflow-y: auto}
        .language-selector {text-align: right; margin-bottom: 10px}
    """) as demo:
        gr.Markdown("# City Hospital - Hospital Booking Assistant")
        gr.Markdown("### മലയാളത്തിലും ഇംഗ്ലീഷിലും സംസാരിക്കുന്ന ആശുപത്രി ബുക്കിംഗ് സഹായി")

        with gr.Row():
            with gr.Column():
                language = gr.Radio(
                    ["English", "Malayalam"],
                    label="Select Language",
                    value="English",
                    interactive=True
                )

        chatbot = gr.Chatbot(
            [],
            elem_id="chatbot",
            label="Chat with Hospital Assistant",
            height=500
        )

        with gr.Row():
            msg = gr.Textbox(
                show_label=False,
                placeholder="Type your message here...",
                container=False
            )
            submit = gr.Button("Send")

        with gr.Row():
            clear = gr.Button("Clear Conversation")

        # Provide instructions
        with gr.Accordion("Instructions", open=False):
            gr.Markdown("""
            ## How to use this hospital booking assistant:

            1. You can chat in English or Malayalam - select your preferred language above.
            2. Ask about doctor availability in different departments.
            3. Check available time slots for appointments.
            4. Book appointments by providing patient details.
            5. Cancel existing appointments if needed.

            ### Example questions in English:
            - Which doctors are available in the cardiology department?
            - I need an appointment with Dr. Priya Nair tomorrow.
            - What are your hospital visiting hours?

            ### Example questions in Malayalam:
            - കാർഡിയോളജി വിഭാഗത്തിൽ ഏതൊക്കെ ഡോക്ടർമാർ ലഭ്യമാണ്?
            - എനിക്ക് നാളെ ഡോ. പ്രിയ നായരുമായി ഒരു അപ്പോയിന്റ്മെന്റ് വേണം.
            - നിങ്ങളുടെ ആശുപത്രി സന്ദർശന സമയങ്ങൾ എന്തൊക്കെയാണ്?
            """)

        chat_history = gr.State([])

        # Set up event handlers
        submit.click(
            process_chat,
            inputs=[msg, chat_history, language, gr.State((model, tokenizer, calendar_service))],
            outputs=[chatbot, chat_history]
        ).then(
            lambda: "",
            None,
            msg
        )

        msg.submit(
            process_chat,
            inputs=[msg, chat_history, language, gr.State((model, tokenizer, calendar_service))],
            outputs=[chatbot, chat_history]
        ).then(
            lambda: "",
            None,
            msg
        )

        clear.click(
            lambda: ([], []),
            inputs=None,
            outputs=[chatbot, chat_history]
        )

        # When language changes, add a system message
        def on_language_change(lang, history):
            if lang == "Malayalam":
                welcome = "ഹലോ, ഞാൻ സിറ്റി ഹോസ്പിറ്റലിന്റെ ഓൺലൈൻ അസിസ്റ്റന്റ് ആണ്. എങ്ങനെ സഹായിക്കാൻ കഴിയും?"
            else:
                welcome = "Hello! I'm the online assistant for City Hospital. How can I help you today?"

            if not history or history[-1][1] != welcome:
                return history + [("", welcome)]
            return history

        language.change(
            on_language_change,
            inputs=[language, chat_history],
            outputs=[chat_history]
        ).then(
            lambda history: (history, history),
            inputs=[chat_history],
            outputs=[chatbot, chat_history]
        )

        # Initial welcome message
        demo.load(
            lambda: ([("", "Hello! I'm the online assistant for City Hospital. How can I help you today?")],
                    [("", "Hello! I'm the online assistant for City Hospital. How can I help you today?")]),
            inputs=None,
            outputs=[chatbot, chat_history]
        )

    return demo

# ===== MAIN EXECUTION =====

def main():
    global HF_TOKEN

    print("===== Malayalam Hospital Booking Chatbot =====")
    print("Using Llama 3.1-8B-Instruct with Google Calendar integration")

    # Install required packages in Colab
    try:
        import IPython
        print("Installing required packages...")
        IPython.get_ipython().system('pip install transformers>=4.37.0')
        IPython.get_ipython().system('pip install accelerate>=0.25.0')
        IPython.get_ipython().system('pip install bitsandbytes>=0.41.0')
        IPython.get_ipython().system('pip install sentencepiece>=0.1.99')
        IPython.get_ipython().system('pip install gradio==3.50.2')
        IPython.get_ipython().system('pip install google-auth google-auth-oauthlib google-auth-httplib2')
        IPython.get_ipython().system('pip install google-api-python-client')
        IPython.get_ipython().system('pip install pytz')
        print("All packages installed successfully!")
    except:
        print("Not running in IPython environment or packages already installed.")

    # Get HF token from user input
    HF_TOKEN = input("Enter your Hugging Face token with access to meta-llama models: ")

    # Load the Llama model and tokenizer
    model, tokenizer = load_llama_model()

    if model is None or tokenizer is None:
        print("Failed to load the model. Please check your Hugging Face token and try again.")
        return

    # Get calendar service
    calendar_service = get_calendar_service()

    # Create and launch the Gradio interface
    demo = create_gradio_interface(model, tokenizer, calendar_service)
    demo.launch(share=True, debug=True)

if __name__ == "__main__":
    main()