from ..base_azure_agent import BaseAzureAgent
from ..azure_connector import AzureConnector
from azure.mgmt.sql.models import Server, Sku
from msrestazure.azure_exceptions import CloudError

class SQLAgent(BaseAzureAgent):
    """
    An agent specialized in handling Azure SQL Database tasks.
    """

    def execute(self, command: str, **kwargs):
        """
        Executes a given command related to Azure SQL Database.
        """
        if command == 'smart_create_sql_database':
            db_name = kwargs.get('database_name', '')
            server_name = kwargs.get('server_name', '')
            resource_group = kwargs.get('resource_group', '')
            print(f"You are about to smart-create an Azure SQL Database '{db_name}' on server '{server_name}' in resource group '{resource_group}'.")
            confirm = input("This will create a new SQL Server (if not exists) and a SQL Database. Are you sure? (yes/no): ")
            if confirm.lower() == 'yes':
                return self._smart_create_sql_database(**kwargs)
            else:
                return {"status": "cancelled", "message": "Smart Create SQL Database command cancelled by user."}
        else:
            raise NotImplementedError(f"Command '{command}' is not supported by SQLAgent.")

    def _smart_create_sql_database(self, subscription_id: str, resource_group: str, location: str, 
                                   server_name: str, database_name: str, admin_username: str, admin_password: str,
                                   sku_name: str = 'Basic', sku_tier: str = 'Basic', sku_capacity: int = 5):
        """
        Creates an Azure SQL Database with sensible defaults.
        Automatically creates a SQL Server if it doesn't exist.
        """
        print(f"SQLAgent: Smart creating SQL Database '{database_name}' on server '{server_name}' in resource group '{resource_group}', location '{location}'...")
        try:
            sql_client = AzureConnector.get_sql_client()

            # 1. Create SQL Server if it doesn't exist
            try:
                sql_client.servers.get(resource_group, server_name)
                print(f"Using existing SQL Server: {server_name}")
            except CloudError:
                print(f"Creating SQL Server: {server_name}")
                server_params = Server(
                    location=location,
                    administrator_login=admin_username,
                    administrator_login_password=admin_password,
                    version='12.0' # SQL Server 2019
                )
                async_server_creation = sql_client.servers.begin_create_or_update(resource_group, server_name, server_params)
                async_server_creation.wait()
                print(f"SQL Server '{server_name}' created.")

            # 2. Create SQL Database
            print(f"Creating SQL Database: {database_name}")
            database_params = SqlDatabase(
                location=location,
                sku=Sku(name=sku_name, tier=sku_tier, capacity=sku_capacity) # Basic tier, 5 DTUs
            )
            async_db_creation = sql_client.databases.begin_create_or_update(resource_group, server_name, database_name, database_params)
            database = async_db_creation.wait()

            return {"status": "success", "message": f"SQL Database '{database_name}' created successfully on server '{server_name}'.", "database_id": database.id}
        except CloudError as e:
            print(f"Azure Error creating SQL Database: {e.message}")
            return {"status": "error", "message": e.message}
        except Exception as e:
            print(f"Error creating SQL Database: {e}")
            return {"status": "error", "message": str(e)}