"""Azure AI model discovery via the Azure Management SDK."""

from __future__ import annotations

import asyncio
import logging
import re
from typing import Any

from azure.identity import DefaultAzureCredential
from azure.mgmt.cognitiveservices import CognitiveServicesManagementClient
from azure.mgmt.resource import ResourceManagementClient
from azure.mgmt.resource.subscriptions import SubscriptionClient

from sage_evaluator.exceptions import DiscoveryError
from sage_evaluator.models import DeployedModel

logger = logging.getLogger(__name__)

_ACCOUNT_NAME_RE = re.compile(r"^[a-zA-Z0-9][a-zA-Z0-9-]{1,62}[a-zA-Z0-9]$")


class AzureModelDiscovery:
    """Discover models deployed in an Azure Cognitive Services account.

    Uses the Azure Management SDK to list deployments.  When only
    ``account_name`` is provided, the subscription and resource group are
    auto-discovered by scanning subscriptions accessible to the credential.

    Args:
        account_name: The Cognitive Services account name, e.g.
            ``"d-ue2-aicorepltfm-aisvcs"``.
        subscription_id: Azure subscription ID.  If *both* this and
            ``resource_group`` are provided, auto-discovery is skipped.
        resource_group: Azure resource group name.

    Example::

        discovery = AzureModelDiscovery("my-account")
        models = await discovery.list_models()
        for m in models:
            print(m.litellm_id, m.capabilities)
    """

    def __init__(
        self,
        account_name: str,
        subscription_id: str | None = None,
        resource_group: str | None = None,
    ) -> None:
        if not _ACCOUNT_NAME_RE.match(account_name):
            raise DiscoveryError(
                f"Invalid account name: {account_name!r}. "
                "Account names must be 3-64 alphanumeric/hyphen characters."
            )
        self.account_name = account_name
        self._subscription_id = subscription_id
        self._resource_group = resource_group
        self._credential = DefaultAzureCredential()

    async def list_models(self) -> list[DeployedModel]:
        """List deployed models in the Azure Cognitive Services account.

        Returns:
            A (possibly empty) list of :class:`DeployedModel` objects.

        Raises:
            DiscoveryError: On authentication or SDK errors.
        """
        logger.info("Discovering models for account %r", self.account_name)
        loop = asyncio.get_running_loop()
        try:
            await loop.run_in_executor(None, self._resolve_account)
        except DiscoveryError:
            raise
        except Exception as exc:
            raise DiscoveryError(f"Failed to resolve account '{self.account_name}': {exc}") from exc

        logger.info("Listing deployments for account %r", self.account_name)
        try:
            deployments = await loop.run_in_executor(None, self._list_deployments)
        except DiscoveryError:
            raise
        except Exception as exc:
            raise DiscoveryError(
                f"Failed to list deployments for '{self.account_name}': {exc}"
            ) from exc

        models = self._map_deployments(deployments)
        logger.info("Discovered %d model(s) for account %r", len(models), self.account_name)
        return models

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _resolve_account(self) -> None:
        """Find subscription_id and resource_group if not already set."""
        if self._subscription_id and self._resource_group:
            logger.info(
                "Using provided subscription=%s, resource_group=%s",
                self._subscription_id,
                self._resource_group,
            )
            return

        logger.info(
            "Auto-discovering subscription and resource group for account %r",
            self.account_name,
        )
        sub_client = SubscriptionClient(credential=self._credential)
        subscriptions = list(sub_client.subscriptions.list())
        if not subscriptions:
            raise DiscoveryError(
                "No Azure subscriptions found for the current credential. "
                "Check that your Azure identity has access to at least one subscription."
            )

        for sub in subscriptions:
            sub_id = sub.subscription_id
            if sub_id is None:
                continue
            rmc = ResourceManagementClient(
                credential=self._credential,
                subscription_id=sub_id,
            )
            resources = list(
                rmc.resources.list(
                    filter=(
                        f"name eq '{self.account_name}' and "
                        "resourceType eq 'Microsoft.CognitiveServices/accounts'"
                    ),
                )
            )
            if resources:
                resource = resources[0]
                resource_id = resource.id
                if resource_id is None:
                    continue
                self._subscription_id = sub_id
                self._resource_group = self._extract_resource_group(resource_id)
                logger.debug(
                    "Resolved account '%s' → subscription=%s, resource_group=%s",
                    self.account_name,
                    self._subscription_id,
                    self._resource_group,
                )
                return

        raise DiscoveryError(
            f"Cognitive Services account '{self.account_name}' not found "
            f"in any of the {len(subscriptions)} accessible subscription(s). "
            "Verify the account name or provide --subscription and --resource-group."
        )

    def _list_deployments(self) -> list[Any]:
        """Call CognitiveServicesManagementClient.deployments.list()."""
        if self._subscription_id is None or self._resource_group is None:
            raise DiscoveryError("_resolve_account must be called before _list_deployments")
        client = CognitiveServicesManagementClient(
            credential=self._credential,
            subscription_id=self._subscription_id,
        )
        return list(
            client.deployments.list(
                resource_group_name=self._resource_group,
                account_name=self.account_name,
            )
        )

    @staticmethod
    def _map_deployments(deployments: list[Any]) -> list[DeployedModel]:
        """Convert Azure SDK Deployment objects to DeployedModel instances."""
        results: list[DeployedModel] = []
        for dep in deployments:
            model = getattr(dep.properties, "model", None)
            if model is None:
                logger.warning("Skipping deployment without model: %s", dep.name)
                continue

            name = getattr(model, "name", "") or ""
            if not name:
                logger.warning("Skipping deployment with empty model name: %s", dep.name)
                continue

            version = getattr(model, "version", "") or ""
            capabilities = AzureModelDiscovery._extract_capabilities(dep)

            results.append(
                DeployedModel(
                    name=name,
                    version=version,
                    capabilities=capabilities,
                    litellm_id=f"azure_ai/{name}",
                )
            )

        logger.debug("Discovered %d model(s) from Azure Management SDK", len(results))
        return results

    @staticmethod
    def _extract_capabilities(deployment: Any) -> list[str]:
        """Extract capabilities from a Deployment object.

        The SDK ``deployment.properties.capabilities`` may be a dict of
        ``{"chatCompletion": "true", ...}`` or ``None``.  Falls back to
        a heuristic based on the model name.
        """
        caps = getattr(deployment.properties, "capabilities", None)
        if caps is not None:
            return AzureModelDiscovery._normalise_capabilities(caps)

        # Heuristic default for known model families
        model = getattr(deployment.properties, "model", None)
        name_lower = (getattr(model, "name", "") or "").lower()
        if any(hint in name_lower for hint in ("gpt", "claude", "llama", "mistral", "phi")):
            return ["chat", "completion"]

        return []

    @staticmethod
    def _normalise_capabilities(caps: Any) -> list[str]:
        """Normalise capability data into a plain list of strings.

        Handles:
        * A dict: ``{"chatCompletion": "true", "embeddings": "false"}``
        * A list of strings: ``["chat", "completion"]``
        * A list of dicts: ``[{"name": "ChatCompletion", "value": "true"}]``
        """
        if isinstance(caps, dict):
            return [k for k, v in caps.items() if str(v).lower() not in ("false", "0", "no")]

        if isinstance(caps, list):
            result: list[str] = []
            for cap in caps:
                if isinstance(cap, str):
                    result.append(cap)
                elif isinstance(cap, dict):
                    name = cap.get("name", "")
                    value = str(cap.get("value", "true")).lower()
                    if name and value not in ("false", "0", "no"):
                        result.append(name.lower())
            return result

        return []

    @staticmethod
    def _extract_resource_group(resource_id: str) -> str:
        """Parse the resource group name from an Azure resource ID string."""
        match = re.search(r"/resourceGroups/([^/]+)", resource_id, re.IGNORECASE)
        if not match:
            raise DiscoveryError(f"Could not parse resource group from resource ID: {resource_id}")
        return match.group(1)
