import os
from sqlalchemy.sql import text
from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker

from semantic_kernel import Kernel
from semantic_kernel.orchestration.sk_context import SKContext
from semantic_kernel.skill_definition import sk_function, sk_function_context_parameter

# Load environment variables from .env file
from dotenv import load_dotenv
load_dotenv()

class SQLGetQuery:
    def __init__(self, kernel: Kernel):
        self._kernel = kernel
        
    @sk_function(
    description="Query SQL database related to user question",
    name="dbQuery",
    )
    @sk_function_context_parameter(name="input", description="SQL query related to user question which was generated by LLM")
    def db_query(self, context: SKContext) -> str:
        # Set SQLAlchemy logging level to WARNING
        username = os.environ['USERNAME']
        password = os.environ['PASSWORD']
        driver = os.environ['DRIVER']
        # schema = context['db_schema']
        server = os.environ['CONNECTION_STRING']
        database = os.environ['DB']
        query = text(context['input'])
        print(f"\n------------------\nSQL QUERY:\n {query}\n------------------\n")
        connection_string = f'mssql+pyodbc://{username}:{password}@{server}/{database}?driver={driver}&TrustServerCertificate=no&Authentication=ActiveDirectoryPassword'
        # Start query
        engine = create_engine(connection_string, echo=False, connect_args={'auto_commit': True}, fast_executemany=True)
        Session = sessionmaker(bind=engine)
        session = Session()
        # Execute the query and fetch the results
        results = session.execute(query).fetchall()
        results = str(results).strip()
        #Update context
        context['queryresult'] = results
        context['db_string'] = server
        context['db_name'] = database
        context['sql_query'] = str(query)
        print(f"\n------------------\nSQL RESULT:\n {results}\n------------------\n")
        
        return results