"""
Data access utilities for HLA-Compass modules.
Architecture: Scoped Vending Machine.

Access is provided via generic interfaces (SQL, Storage) but scoped to specific
Provider/Catalog domains to ensure extensibility and isolation.
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)


def _normalize_hla_locus(locus: str | None) -> str | None:
    if locus is None:
        return None
    normalized = str(locus).strip().upper()
    if not normalized:
        return None
    if not normalized.startswith("HLA-"):
        normalized = f"HLA-{normalized}"
    return normalized


def _normalize_hla_resolution(resolution: str) -> str:
    normalized = str(resolution or "2-digit").strip().lower()
    if normalized not in {"2-digit", "4-digit", "full"}:
        raise DataAccessError("Unsupported HLA resolution")
    return normalized


def _format_hla_allele_resolution(allele_name: str, resolution: str) -> str:
    if resolution == "full" or "*" not in allele_name:
        return allele_name

    prefix, suffix = allele_name.split("*", 1)
    segments = [segment for segment in suffix.split(":") if segment]
    if not segments:
        return allele_name
    if resolution == "2-digit":
        return f"{prefix}*{segments[0]}"
    return f"{prefix}*{':'.join(segments[:2])}" if len(segments) >= 2 else f"{prefix}*{segments[0]}"


def _get_api_client(api_client):
    if api_client is not None:
        return api_client

    from .client import APIClient

    return APIClient()


class DataAccessError(Exception):
    """Error accessing HLA-Compass data"""

    pass


class DataClient:
    """
    Entry point for accessing scientific data from a specific catalog.

    Example:
        >>> genetics = DataClient(provider="alithea", catalog="genetics")
        >>> genetics.sql.query("SELECT * FROM genes")
    """

    def __init__(
        self,
        provider: str = "alithea-bio",
        catalog: str = "immunopeptidomics",
        api_client=None,
    ):
        self.provider = provider
        self.catalog = catalog
        self.api = api_client

        # Initialize sub-clients
        self.sql = SQLClient(self)
        # Backwards-compatible aliases:
        # - `storage` is the ergonomic name used by most examples.
        # - `catalog_storage` is the explicit name used internally/tests.
        self.storage = StorageClient(self)
        self.catalog_storage = self.storage

    def _get_api(self):
        """Lazy load API client if not provided"""
        if self.api:
            return self.api
        # Import here to avoid circular dependency if client imports data
        from .client import APIClient

        self.api = APIClient()
        return self.api

    def _resolve_endpoint_prefix(self) -> str:
        """Return /v1/api/data for API-key auth, /v1/data for user sessions."""
        try:
            api = self._get_api()
            headers = api.auth.get_headers() if getattr(api, "auth", None) else {}
            if any(str(k).lower() == "x-api-key" for k in (headers or {}).keys()):
                return "/v1/api/data"
        except Exception:
            pass
        return "/v1/data"


class SQLClient:
    """
    Execute SQL queries against the Catalog's Schema.

    For available tables and columns, see the Data Catalog documentation:
    https://docs.alithea.bio/sdk-reference/guides/data-access
    """

    def __init__(self, parent: DataClient):
        self.parent = parent
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

    def query(self, sql: str, params: Optional[List[Any]] = None) -> Dict[str, Any]:
        """
        Execute a raw SQL query.

        Args:
            sql: SQL query string with %s placeholders for parameters
            params: List of parameter values to bind safely

        Returns:
            Dictionary with 'columns', 'data', 'count'

        Example:
            >>> result = self.data.sql.query(
            ...     "SELECT sequence, mass FROM peptides WHERE length >= %s LIMIT 10",
            ...     params=[8]
            ... )
            >>> print(result["data"])
        """
        api = self.parent._get_api()

        try:
            payload: dict[str, Any] = {"sql": sql}
            if params:
                payload["params"] = params

            prefix = self.parent._resolve_endpoint_prefix()
            endpoint = f"{prefix}/{self.parent.provider}/{self.parent.catalog}/query"
            result = api._make_request("POST", endpoint, json_data=payload)
            return result

        except Exception as e:
            self.logger.error(f"SQL query failed: {e}")
            raise DataAccessError(f"Failed to execute query: {str(e)}") from e

    def query_df(self, sql: str, params: Optional[List[Any]] = None, engine: str = "pandas"):
        """Execute a SQL query and return results as a DataFrame.

        Args:
            sql: SQL query string with %s placeholders for parameters
            params: List of parameter values to bind safely
            engine: ``"pandas"`` or ``"polars"``

        Returns:
            pandas.DataFrame or polars.DataFrame

        Raises:
            ImportError: If the requested engine is not installed
            DataAccessError: If the query fails

        Example:
            >>> df = self.data.sql.query_df(
            ...     "SELECT sequence, mass FROM peptides WHERE length >= %s LIMIT 10",
            ...     params=[8],
            ... )
            >>> print(df.head())
        """
        result = self.query(sql, params)
        columns = result.get("columns", [])
        data = result.get("data", [])

        if engine == "pandas":
            try:
                import pandas as pd
            except ImportError:
                raise ImportError(
                    "pandas is required for query_df(). Install with: pip install 'hla-compass[data]'"
                ) from None
            return pd.DataFrame(data, columns=columns if columns else None)

        if engine == "polars":
            try:
                import polars as pl
            except ImportError:
                raise ImportError(
                    "polars is required for query_df(engine='polars'). Install with: pip install 'hla-compass[data]'"
                ) from None
            if columns:
                return pl.DataFrame(data, schema=columns, orient="row")
            return pl.DataFrame(data)

        raise ValueError(f"Unknown engine: {engine!r}. Use 'pandas' or 'polars'.")


class StorageClient:
    """
    Access the Catalog's Object Storage.
    """

    def __init__(self, parent: DataClient):
        self.parent = parent
        self.logger = logging.getLogger(f"{__name__}.{self.__class__.__name__}")

    def get_credentials(self, mode: str = "read") -> Dict[str, Any]:
        """
        Get temporary AWS STS credentials scoped to the catalog's bucket prefix.
        """
        api = self.parent._get_api()

        try:
            prefix = self.parent._resolve_endpoint_prefix()
            endpoint = f"{prefix}/{self.parent.provider}/{self.parent.catalog}/storage/token"
            creds = api._make_request("GET", endpoint, params={"mode": mode})
            return creds
        except Exception as e:
            self.logger.error(f"Failed to get storage credentials: {e}")
            raise DataAccessError(f"Storage access denied: {str(e)}") from e

    def get_s3_client(self, mode: str = "read"):
        """
        Get a boto3 S3 client configured with scoped credentials.
        """
        try:
            import boto3
        except ImportError:
            raise ImportError("boto3 is required for get_s3_client") from None

        creds = self.get_credentials(mode)
        return boto3.client(
            "s3",
            aws_access_key_id=creds["accessKeyId"],
            aws_secret_access_key=creds["secretAccessKey"],
            aws_session_token=creds["sessionToken"],
            region_name=creds["region"],
        )

    def read_parquet(self, key: str, engine: str = "polars"):
        """
        Read a parquet file from the catalog's storage.

        Args:
            key: Relative path (e.g. "runs/123.parquet")
            engine: "polars" or "pandas"
        """
        creds = self.get_credentials(mode="read")

        # Construct full S3 URI using the bucket and prefix returned by the token vendor
        bucket = creds["bucket"]
        prefix = creds.get("prefix", "")

        # Handle if key is already absolute
        if key.startswith("s3://"):
            uri = key
        else:
            # Ensure cleanly joined path
            clean_prefix = prefix.rstrip("/")
            clean_key = key.lstrip("/")
            if clean_prefix:
                uri = f"s3://{bucket}/{clean_prefix}/{clean_key}"
            else:
                uri = f"s3://{bucket}/{clean_key}"

        storage_options = {
            "key": creds["accessKeyId"],
            "secret": creds["secretAccessKey"],
            "token": creds["sessionToken"],
            "region": creds["region"],
        }

        if engine == "polars":
            try:
                import polars as pl

                return pl.read_parquet(uri, storage_options=storage_options)
            except ImportError:
                raise ImportError("polars is required") from None

        elif engine == "pandas":
            try:
                import pandas as pd

                return pd.read_parquet(uri, storage_options=storage_options)
            except ImportError:
                raise ImportError("pandas is required") from None
        else:
            raise ValueError(f"Unknown engine: {engine}")


class PeptideData:
    """Convenience wrapper for peptide-centric access patterns.

    Prefer direct database access when a `db_client` (e.g. `ScientificQuery`) is provided;
    otherwise fall back to the public API client.
    """

    def __init__(self, api_client=None, db_client=None):
        self.api = api_client
        self.db = db_client

    def _get_api(self):
        self.api = _get_api_client(self.api)
        return self.api

    @classmethod
    def from_module(cls, module) -> "PeptideData":
        """Create PeptideData from a Module instance."""
        return cls(api_client=module.api, db_client=module.db)

    def search(self, *, sequence: str | None = None, limit: int = 100, **filters: Any) -> List[Dict[str, Any]]:
        if self.db is not None and sequence is not None:
            return self.db.execute_function(
                "search_peptides_by_sequence",
                {"pattern": sequence, "max_results": limit},
            )

        request_filters: Dict[str, Any] = {}
        if sequence is not None:
            request_filters["sequence"] = sequence
        request_filters.update(filters)
        return self._get_api().get_peptides(filters=request_filters or None, limit=limit)

    def search_by_hla(
        self,
        allele_name: str,
        *,
        binding_score_min: float = 0.0,
        limit: int = 100,
        **filters: Any,
    ) -> List[Dict[str, Any]]:
        if self.db is not None:
            payload = {
                "allele_name": allele_name,
                "min_intensity": binding_score_min,
                "max_results": limit,
            }
            payload.update(filters)
            return self.db.execute_function("search_peptides_by_hla", payload)

        raise DataAccessError(
            "Peptide HLA search is not exposed through the public REST data endpoints. "
            "Use module-runtime database access or self.data.sql.query(...) instead."
        )


class ProteinData:
    """Convenience wrapper for protein-centric access patterns.

    Prefer direct database access when a ``db_client`` is provided;
    otherwise fall back to the public API client.
    """

    def __init__(self, api_client=None, db_client=None):
        self.api = api_client
        self.db = db_client

    def _get_api(self):
        self.api = _get_api_client(self.api)
        return self.api

    @classmethod
    def from_module(cls, module) -> "ProteinData":
        """Create ProteinData from a Module instance."""
        return cls(api_client=module.api, db_client=module.db)

    def search(
        self,
        *,
        accession: str | None = None,
        gene_name: str | None = None,
        limit: int = 100,
        **filters: Any,
    ) -> List[Dict[str, Any]]:
        request_filters: Dict[str, Any] = {}
        if accession is not None:
            request_filters["accession"] = accession
        if gene_name is not None:
            request_filters["gene_name"] = gene_name
        request_filters.update(filters)
        return self._get_api().get_proteins(filters=request_filters or None, limit=limit)

    def get(self, protein_id: str) -> Dict[str, Any]:
        return self._get_api().get_protein(protein_id)


class SampleData:
    """Convenience wrapper for sample-centric access patterns.

    Prefer direct database access when a ``db_client`` is provided;
    otherwise fall back to the public API client.
    """

    def __init__(self, api_client=None, db_client=None):
        self.api = api_client
        self.db = db_client

    def _get_api(self):
        self.api = _get_api_client(self.api)
        return self.api

    @classmethod
    def from_module(cls, module) -> "SampleData":
        """Create SampleData from a Module instance."""
        return cls(api_client=module.api, db_client=module.db)

    def search(
        self,
        *,
        tissue: str | None = None,
        disease: str | None = None,
        limit: int = 100,
        **filters: Any,
    ) -> List[Dict[str, Any]]:
        request_filters: Dict[str, Any] = {}
        if tissue is not None:
            request_filters["tissue"] = tissue
        if disease is not None:
            request_filters["disease"] = disease
        request_filters.update(filters)
        return self._get_api().get_samples(filters=request_filters or None, limit=limit)

    def get(self, sample_id: str) -> Dict[str, Any]:
        return self._get_api().get_sample(sample_id)


class HLAData:
    """Convenience wrapper for HLA allele access patterns."""

    def __init__(self, api_client=None, db_client=None):
        self.api = api_client
        self.db = db_client

    def _get_api(self):
        self.api = _get_api_client(self.api)
        return self.api

    @classmethod
    def from_module(cls, module) -> "HLAData":
        """Create HLAData from a Module instance."""
        return cls(api_client=module.api, db_client=module.db)

    def search_alleles(
        self,
        *,
        locus: str | None = None,
        resolution: str = "2-digit",
        limit: int = 100,
    ) -> List[str]:
        if self.db is not None:
            normalized_locus = _normalize_hla_locus(locus)
            normalized_resolution = _normalize_hla_resolution(resolution)

            sql = "SELECT allele_name FROM scientific.hla_alleles"
            params: list[Any] = []
            if normalized_locus:
                sql += " WHERE gene = %s"
                params.append(normalized_locus)
            sql += " ORDER BY allele_name ASC"

            rows = self.db.execute_readonly(sql, params or None)
            alleles: list[str] = []
            seen: set[str] = set()
            for row in rows:
                allele_name = str(row.get("allele_name") or "").strip()
                if not allele_name:
                    continue
                resolved = _format_hla_allele_resolution(allele_name, normalized_resolution)
                if resolved not in seen:
                    seen.add(resolved)
                    alleles.append(resolved)
            return alleles[: max(int(limit), 0)]

        results = self._get_api().get_hla_alleles(locus=locus, resolution=resolution)
        return results[: max(int(limit), 0)]
