from langchain.tools import BaseTool
from langchain_openai import ChatOpenAI
from langchain.agents import AgentExecutor, create_openai_tools_agent
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_community.utilities import SQLDatabase
from langchain.schema import SystemMessage, HumanMessage, AIMessage
import os
import psycopg2
import gradio as gr

# Set up database connection details
pg_uri = "postgresql://ridhima:0skESLQ9D6c3m7smqwG47peapk7HzVvu@dpg-cq2h613v2p9s73esp8eg-a.singapore-postgres.render.com/hr_qugd"

def get_db_connection():
    conn = psycopg2.connect(pg_uri)
    return conn

class LeaveRequestInfoTool(BaseTool):
    name = "leave_request_information"
    description = "Provides information about the leave request process and database schema."

    def _run(self, query: str) -> str:
        return """
        Leave Request Process:
        1. Collect employee ID
        2. Ask for leave type
        3. Get reason for leave
        4. Get start date of leave
        5. Get end date of leave
        6. Calculate duration (in days)
        7. Insert data into request table
        8. Confirm submission to user

        Database Schema:
        Table: request
        Columns:
        - leave_id (auto-increment integer)
        - employee_id (integer)
        - leave_type (text)
        - reason (text)
        - start_of_leave (date)
        - end_of_leave (date)
        - duration (integer, calculated in days)
        - leave_status (text, default 'Pending')

        Instructions:
        - Collect all necessary information from the user one by one.
        - Calculate the duration as the number of days between start_of_leave and end_of_leave.
        - Once all information is collected, formulate an SQL INSERT statement for the 'request' table.
        - REMEMBER TO EXECUTE THE INSERT QUERY FOR EACH REQUEST ONLY ONCE.
        - CRITICAL: EXECUTE THE INSERT QUERY FOR EACH REQUEST ONLY ONCE.
        - If you receive an error message saying an insertion has already been made, DO NOT attempt to insert again.
        - Instead, inform the user that their request has been submitted and ask if they need anything else.
        - After insertion, confirm to the user that their request has been submitted.
        """

class SQLAgentTool(BaseTool):
    name = "sql_agent"
    description = "Use this tool to interact with the database and execute SQL queries."
    
    def _run(self, query: str) -> str:
        print("Executing Query:  ", query)
        conn = get_db_connection()
        cur = conn.cursor()
        try:
            cur.execute(query)
            conn.commit()
            result = "Your leave request has been submitted successfully."
        except Exception as e:
            conn.rollback()
            result = f"An error occurred: {str(e)}"
        finally:
            cur.close()
            conn.close()
        return result

prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content="""You are an HR assistant. You can help with leave requests and provide information about company policies.
    - For leave requests, ask for each piece of information one at a time. After collecting all information, use the sql_agent tool to INSERT the data into the 'request' table.
    - For information queries about company policies, use the rag_info tool to provide accurate information from the RAG-trained model.
    
    The pieces of information you need to collect for leave requests are:
    1. Employee ID
    2. Leave Type
    3. Reason for Leave
    4. Start Date of Leave (YYYY-MM-DD)
    5. End Date of Leave (YYYY-MM-DD)
    Calculate the duration as the number of days between start and end dates using PostgreSQL functions and convert duration to integer.
    Use 'Pending' as the default leave_status.

    Remember to use the correct column names as per the table structure:
    request(leave_id, employee_id, leave_type, reason, start_of_leave, end_of_leave, duration, leave_status)
    Where leave_id is auto-increment and should not be included in the INSERT statement.
    
    For company policies, provide detailed and accurate information based on the RAG-trained model.
    """),
    MessagesPlaceholder(variable_name="chat_history"),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])