# ruff: noqa: T201
"""Configuration and credentials import handler for Azure Foundry."""

from azure.identity import DefaultAzureCredential
from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient
from azure.mgmt.subscription import SubscriptionClient

from yera.config2.dataclasses import CredentialsMap, LLMConfig, ModelRegistry
from yera.config2.keyring import DevKeyring
from yera.config2.setup_handlers.base import BaseProviderSetup


def _discover_azure_openai_accounts(credential: DefaultAzureCredential) -> list[dict]:
    subscription_client = SubscriptionClient(credential)

    accounts = []

    for subscription in subscription_client.subscriptions.list():
        sub_id = subscription.subscription_id

        cog_client = CognitiveServicesManagementClient(credential, sub_id)

        for account in cog_client.accounts.list():
            if account.kind == "OpenAI":
                resource_group = account.id.split("/")[4]

                accounts.append(
                    {
                        "resource_name": account.name,
                        "endpoint": account.properties.endpoint,
                        "subscription_id": sub_id,
                        "resource_group": resource_group,
                    }
                )

    return accounts


def _list_deployments(credential: DefaultAzureCredential, account: dict) -> list[dict]:
    client = CognitiveServicesManagementClient(credential, account["subscription_id"])

    deployments = client.deployments.list(
        resource_group_name=account["resource_group"],
        account_name=account["resource_name"],
    )

    return [
        {
            "name": d.name,
            "model": d.properties.model.name,
            "version": d.properties.model.version,
        }
        for d in deployments
    ]


def _get_api_key(credential: DefaultAzureCredential, account: dict) -> str:
    client = CognitiveServicesManagementClient(credential, account["subscription_id"])

    keys = client.accounts.list_keys(
        resource_group_name=account["resource_group"],
        account_name=account["resource_name"],
    )

    return keys.key1


class AzureSetup(BaseProviderSetup):
    """Setup handler for Azure OpenAI credentials and deployments.

    Manages automatic discovery and manual configuration of Azure OpenAI
    resources. Supports automatic credential detection via Azure CLI login
    (DefaultAzureCredential) and manual entry fallback.
    """

    def __init__(self):
        """Initialise the Azure OpenAI setup handler."""
        super().__init__("azure")

    def detect_creds(self) -> dict[str, str] | None:
        """Automatically detect Azure OpenAI credentials via Azure CLI.

        Uses DefaultAzureCredential to authenticate and discover Azure OpenAI
        accounts across all accessible subscriptions. Automatically retrieves
        API keys for discovered accounts. Requires you to log in with the cli
        beforehand.

        Returns:
            Dictionary containing resource_name, endpoint, subscription_id,
            resource_group, and api_key if a single account is found.
            None if no accounts are found or authentication fails.

        Raises:
            RuntimeError: If multiple Azure OpenAI accounts are found.

        Note:
            Requires user to be logged in via 'az login' and have appropriate
            permissions to list accounts and retrieve API keys.
        """
        credential = DefaultAzureCredential()
        accounts = _discover_azure_openai_accounts(credential=credential)

        if len(accounts) == 0:
            return None

        if len(accounts) > 1:
            print("Found multiple Azure OpenAI accounts:")
            for i, acc in enumerate(accounts):
                acc_name = acc["resource_name"]
                subsc_id = acc["subscription_id"]
                print(f"  [{i + 1}] {acc_name} (subscription: {subsc_id[:8]}...)")
            choice = int(input("Select account number: ")) - 1
            account = accounts[choice]
        else:
            account = accounts[0]

        # Add API version detection/configuration:
        account["api_key"] = _get_api_key(credential=credential, account=account)
        return account

    def ask_for_creds(self) -> dict[str, str] | None:
        """Prompt user to manually enter Azure OpenAI credentials.

        Interactively requests all required Azure OpenAI configuration
        including account name, endpoint, subscription ID, resource group,
        and API key.

        Returns:
            Dictionary with resource_name, endpoint, subscription_id,
            resource_group, and api_key if user provides credentials.
            None if user declines to enter credentials.
        """
        ask = self._confirm_ask_for_creds()
        if not ask:
            return None

        fields = [
            "resource_name",
            "resource_group",
            "endpoint",
            "subscription_id",
            "api_key",
        ]
        delimiter = "\n * "
        field_list = delimiter.join(fields)
        print(
            f"To get everything set up we will need some of "
            f"your Azure details:{delimiter}{field_list}"
        )
        print("resource_name is your Microsoft Azure Resource Name")

        account = {}

        for field in fields:
            print("Please enter your " + field + ":")
            account[field] = input(" > ")

        return account

    def validate_creds(self, creds: dict[str, str]) -> None:
        """Validate Azure OpenAI credential format and completeness.

        Checks that all required fields are present and validates format
        of endpoint URL, API key length, and subscription ID structure.

        Args:
            creds: Dictionary containing Azure OpenAI credentials.

        Raises:
            ValueError: If required fields are missing, endpoint format is
                invalid, or API key appears incomplete.
        """
        # Check required fields are present
        required_fields = [
            "resource_name",
            "endpoint",
            "subscription_id",
            "resource_group",
            "api_key",
        ]
        missing = [f for f in required_fields if f not in creds or not creds[f]]

        if missing:
            raise ValueError(f"Missing required fields: {', '.join(missing)}")

        # Validate endpoint format
        endpoint = creds["endpoint"]
        if not endpoint.startswith("https://"):
            raise ValueError("Endpoint must start with 'https://'")

        if not (
            ".openai.azure.com" in endpoint
            or ".api.cognitive.microsoft.com" in endpoint
        ):
            print("  Warning: Endpoint doesn't match expected Azure OpenAI format")

        # Validate API key format (basic check)
        api_key = creds["api_key"]
        if len(api_key) < 20:
            raise ValueError(
                "API key appears too short. Please ensure you've copied the complete key."
            )

        # Validate subscription_id format (should be a GUID)
        subscription_id = creds["subscription_id"]
        if len(subscription_id) != 36 or subscription_id.count("-") != 4:
            print("  Warning: Subscription ID doesn't look like a valid GUID")

    def fetch_models(self, creds_map: CredentialsMap) -> ModelRegistry:
        """Fetch available model deployments from Azure OpenAI account.

        Uses Azure Management SDK to list all model deployments in the
        configured Azure OpenAI account and creates LLMConfig entries
        for each deployment.

        Args:
            creds_map: Credentials map containing Azure credential labels.

        Returns:
            ModelRegistry containing LLMConfig entries for all discovered
            model deployments.

        Note:
            Requires DefaultAzureCredential authentication to access the
            Azure Management API.
        """
        account = {
            k: DevKeyring.get(f"providers.{v}")
            for k, v in creds_map.providers["azure"].items()
        }

        credential = DefaultAzureCredential()
        deployments = _list_deployments(credential, account)

        def _make_model_config(d: dict) -> LLMConfig:
            return LLMConfig(
                id=f"azure.{d['name']}".lower(),
                display_name=d["model"],
                credentials="azure",
                provider="azure",
                interface="azure-sdk",
                model_id=d["name"],
            )

        model_configs = [_make_model_config(d) for d in deployments]
        names = "  \n   * ".join(m.display_name for m in model_configs)
        print(f"  Adding {len(model_configs)} LLMs from Azure:\n   * {names}")

        return ModelRegistry(
            llm={m.id: m for m in model_configs},
        )
