"""
Azure Blob Storage source connector.

Auth (env vars):
    AZURE_STORAGE_CONNECTION_STRING              (all-in-one)
    AZURE_STORAGE_ACCOUNT + AZURE_STORAGE_KEY   (explicit)
    AZURE_STORAGE_ACCOUNT alone                  (uses DefaultAzureCredential — managed identity,
                                                  service principal, CLI)

URI format:
    az://container/prefix/
    az://container/docs/report.pdf
    abfs://container/prefix/  (ADLS Gen2 alias)
"""

from __future__ import annotations

import contextlib
import fnmatch
import os
import tempfile
from collections.abc import Iterator
from pathlib import Path
from urllib.parse import urlparse

from .base import SourceFile
from .registry import register_source


class AzureBlobSource:
    scheme = "az"

    def _get_service_client(self):
        try:
            from azure.storage.blob import BlobServiceClient
        except ImportError as e:
            raise ImportError(
                "azure-storage-blob not installed. Run: pip install 'coresdk-extractor[azure]'"
            ) from e

        conn_str = os.environ.get("AZURE_STORAGE_CONNECTION_STRING")
        if conn_str:
            return BlobServiceClient.from_connection_string(conn_str)

        account = os.environ.get("AZURE_STORAGE_ACCOUNT")
        key = os.environ.get("AZURE_STORAGE_KEY")
        if account and key:
            return BlobServiceClient(
                account_url=f"https://{account}.blob.core.windows.net",
                credential=key,
            )
        if account:
            try:
                from azure.identity import DefaultAzureCredential

                credential = DefaultAzureCredential()
            except ImportError as e:
                raise ImportError(
                    "azure-identity not installed. Run: pip install 'coresdk-extractor[azure]'"
                ) from e
            return BlobServiceClient(
                account_url=f"https://{account}.blob.core.windows.net",
                credential=credential,
            )
        raise ValueError(
            "Azure auth not configured. Set AZURE_STORAGE_CONNECTION_STRING or "
            "AZURE_STORAGE_ACCOUNT + AZURE_STORAGE_KEY (or use managed identity)."
        )

    def list(self, uri: str, glob_pattern: str | None = None) -> Iterator[SourceFile]:
        parsed = urlparse(uri)
        container = parsed.netloc
        prefix = parsed.path.lstrip("/")
        svc = self._get_service_client()
        container_client = svc.get_container_client(container)
        for blob in container_client.list_blobs(name_starts_with=prefix):
            filename = blob.name.split("/")[-1]
            if not filename:
                continue
            if glob_pattern and not fnmatch.fnmatch(filename, glob_pattern):
                continue
            yield SourceFile(
                uri=f"az://{container}/{blob.name}",
                filename=filename,
                size_bytes=blob.size,
                etag=blob.etag,
                last_modified=blob.last_modified.isoformat() if blob.last_modified else None,
                content_type=blob.content_settings.content_type if blob.content_settings else None,
            )

    @contextlib.contextmanager
    def fetch(self, source_file: SourceFile, tmp_dir: Path | None = None) -> Iterator[Path]:
        parsed = urlparse(source_file.uri)
        container = parsed.netloc
        blob_name = parsed.path.lstrip("/")
        svc = self._get_service_client()
        blob_client = svc.get_blob_client(container=container, blob=blob_name)
        suffix = Path(source_file.filename).suffix or ".bin"
        with tempfile.NamedTemporaryFile(suffix=suffix, dir=tmp_dir, delete=False) as tmp:
            tmp_path = Path(tmp.name)
        try:
            with open(tmp_path, "wb") as f:
                data = blob_client.download_blob()
                data.readinto(f)
            yield tmp_path
        finally:
            with contextlib.suppress(FileNotFoundError):
                tmp_path.unlink()


class ADLSSource(AzureBlobSource):
    scheme = "abfs"


register_source(AzureBlobSource())
register_source(ADLSSource())
