import json
from fastapi import FastAPI
from pydantic import BaseModel
import mysql.connector
from fastapi import APIRouter
import mysql.connector
from azure.identity import ClientSecretCredential
from azure.keyvault.secrets import SecretClient
from mysql.connector import pooling
import json
from azure.storage.blob import BlobServiceClient
from azure.storage.blob import BlobPrefix
import pyodbc
import agenthink.utils as utils
import os
import json
import dotenv
import logging
import re
dotenv.load_dotenv()

# Create and configure logger
logging.basicConfig(filename="datastore_library.log",
                    format='%(asctime)s %(message)s',
                    filemode='w')



class DBConnector:
    _class_cache = {}   # session_id -> instance

    def __init__(self, session_id: str, user_id: str, workflow_id: str):
        self.session_id = session_id
        self._initialized = False
        self._workflow_id = None
        self._user_id = None
        self.connection_object_dict = {}


        self.__account_name = os.getenv("AZURE_STORAGE_DATASETS_ACCOUNT_NAME")
        self.__account_key = os.getenv("AZURE_STORAGE_DATASETS_ACCOUNT_KEY")
        self.__container_name = os.getenv("AZURE_STORAGE_DATASETS_CONTAINER_NAME")
        connection_str = (
            f"DefaultEndpointsProtocol=https;"
            f"AccountName={self.__account_name};"
            f"AccountKey={self.__account_key};"
            f"EndpointSuffix=core.windows.net")
        
        blob_service = BlobServiceClient.from_connection_string(connection_str)
        container_client = blob_service.get_container_client(container_name)
        json_data = json.loads(data.decode("utf-8"))
        self.no_of_datastores = len(json_data)
        secret_names = [entry["key"] for entry in json_data]

        self.__client_id = os.getenv("CLIENT_ID")
        self.__tenant_id = os.getenv("TENANT_ID")    
        self.__client_secret = os.getenv("CLIENT_SECRET")
        self.__vault_url = os.getenv("VAULT_URL")
        
        

        for dict_creds in json_data:
            cleaned_key = self.sanitize_secret_name(dict_creds["key"])
            datastore_type = dict_creds["datastore_type"]
            name = dict_creds["name"]

            credential = ClientSecretCredential(
                tenant_id=self.__tenant_id,
                client_id=self.__client_id,
                client_secret=self.__client_secret
            )

            client = SecretClient(vault_url=self.__vault_url, credential=credential)
            secret_value = client.get_secret(cleaned_key)
            secret = json.loads(secret_value.value)

            # MySQL connection
            if datastore_type == "mysql":
                connection_object = self.__connect_mysql(secret)
                try:
                    self.connection_object_dict[f"{secret["database_name"]}"] = connection_object
                except Exception as e:
                    print(f"Unable to create connection object for database name: {secret["database_name"]}")
                pass

            if datastore_type == "mssql":
                connection_object = self.__connect_mssql(secret)
                self.connection_object_dict[f"{secret["database_name"]}"] = connection_object
                pass




        

    # ---- class method to create instance ----
    @classmethod
    def get(cls, session_id: str, user_id: str, workflow_id: str):
        if session_id not in cls._class_cache:
            cls._class_cache[session_id] = DBConnector(session_id,user_id,workflow_id)
        return cls._class_cache[session_id]

    # ---- public method ----
    def execute(self, query: str, workflow_id: str, user_id: str):
        if not self._initialized:
            self._workflow_id = workflow_id
            self._user_id = user_id

        return {
            "session": self.session_id,
            "workflow": self._workflow_id,
            "user": self._user_id,
            "query": query
        }
    
    def __connect_mysql(self, secret):

        config = {
            "host": secret["cluster_ip"],
            "port": secret["port"],
            "user": secret["username"],
            "password": secret["password"],
            "database": secret["database_name"]
        }

        try:
        # self.connect_mysql(config)
            conn = mysql.connector.connect(**config)
            print(f"MY SQL connection successfully established")
            return conn
        
        except Exception as e:
            print("MY SQL connection failed:", e)


    def __connect_mssql(self, secret):
        connection_string = (
            f"mssql+pyodbc://{secret["username"]}:{secret["password"]}@{secret["cluster_ip"]}:{secret["port"]}/{secret["database_name"]}"
            "?driver=ODBC+Driver+17+for+SQL+Server&Encrypt=no"
        )
        try:
            conn = pyodbc.connect(connection_string)
            cursor = conn.cursor()
            row = cursor.fetchone()
            print("Connected! SQL version:", row[0])
        except Exception as e:
            print("MS SQL Connection failed:", e)



    def sanitize_secret_name(self, secret_name: str) -> str:
        """
        Convert secret name to a valid Azure Key Vault format:
        - Replaces spaces and underscores with hyphens
        - Removes any characters not allowed in Key Vault secret names
        """
        # Replace spaces and underscores with hyphens
        sanitized = re.sub(r"[ _]+", "-", secret_name)
        
        # Remove any characters that are not letters, numbers, or hyphens
        sanitized = re.sub(r"[^a-zA-Z0-9\-]", "", sanitized)
        
        return sanitized

