import os
from dotenv import load_dotenv
load_dotenv()
from langchain_huggingface import HuggingFaceEndpoint
from langchain_community.utilities.sql_database import SQLDatabase
from langchain_community.agent_toolkits.sql.base import SQLDatabaseToolkit, create_sql_agent
from langchain.agents.agent_types import AgentType

# Set your Hugging Face API key and model
HUGGINGFACEHUB_API_TOKEN = os.getenv("HUGGINGFACEHUB_API_TOKEN")
HUGGINGFACE_MODEL = os.getenv("HUGGINGFACE_MODEL", "microsoft/phi-2")

# Set your Postgres connection details
PG_HOST = os.getenv("PG_HOST", "localhost")
PG_PORT = os.getenv("PG_PORT", "5432")
PG_USER = os.getenv("PG_USER", "postgres")
PG_PASSWORD = os.getenv("PG_PASSWORD", "password")
PG_DATABASE = os.getenv("PG_DATABASE", "postgres")


# Create the SQLAlchemy connection string
connection_string = f"postgresql+psycopg2://{PG_USER}:{PG_PASSWORD}@{PG_HOST}:{PG_PORT}/{PG_DATABASE}"

def get_sql_agent(verbose=False):
    db = SQLDatabase.from_uri(connection_string)
    llm = HuggingFaceEndpoint(
        repo_id=HUGGINGFACE_MODEL,
        huggingfacehub_api_token=HUGGINGFACEHUB_API_TOKEN,
        temperature=0.2,
        max_new_tokens=256
    )
    agent_executor = create_sql_agent(
        llm=llm,
        db=db,
        agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        verbose=verbose
    )
    return agent_executor