# Copyright (c) 2026, Alexander A. Alsalihi, Robert M. Flight,
# Hunter N.B. Moseley. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#
# * Redistributions of source code must retain the above copyright notice, this
#   list of conditions and the following disclaimer.
# * Redistributions in binary form must reproduce the above copyright notice,
#   this list of conditions and the following disclaimer in the documentation
#   and/or other materials provided with the distribution.
# * All advertising materials mentioning features or use of this software must
#   display the following acknowledgement: This product includes software
#   developed by the copyright holder.
# * Neither the name of the copyright holder nor the names of its contributors
#   may be used to endorse or promote products derived from this software without
#   specific prior written permission.
# * If the source code is used in a published work, then proper citation of the
#   source code must be included with the published work.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS AS IS AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
"""Resource descriptions and duffel URL-path construction.

This module defines a small set of resource description dataclasses for 
recount3. A description is a validated, immutable-ish bundle of parameters
(organism, project, etc.) that can deterministically construct the relative
path to a resource in the repository.

The main entry point is :class:`R3ResourceDescription`, which acts as a
multi-factory: instantiating ``R3ResourceDescription(resource_type=...)``
returns an instance of the registered concrete subclass for that
``resource_type``.

Registered resource types
  ==========================================  ============================
  ``resource_type`` string                    Class
  ==========================================  ============================
  ``"annotations"``                           :class:`R3Annotations`
  ``"count_files_gene_or_exon"``              :class:`R3GeneOrExonCounts`
  ``"count_files_junctions"``                 :class:`R3JunctionCounts`
  ``"metadata_files"``                        :class:`R3ProjectMetadata`
  ``"bigwig_files"``                          :class:`R3BigWig`
  ``"data_sources"``                          :class:`R3DataSources`
  ``"data_source_metadata"``                  :class:`R3DataSourceMetadata`
  ==========================================  ============================

Valid values for the ``organism`` and ``data_source`` fields are exposed as
the module-level constants :data:`VALID_ORGANISMS` and
:data:`VALID_DATA_SOURCES`.

Typical usage::

  from recount3._descriptions import R3ResourceDescription

  desc = R3ResourceDescription(
      resource_type="count_files_gene_or_exon",
      organism="human",
      data_source="sra",
      genomic_unit="gene",
      project="SRP107565",
      annotation_extension="G026",
  )
  path = desc.url_path()

  # Resource type can also be passed as the first positional argument:
  desc = R3ResourceDescription("data_sources", organism="human")
"""

from __future__ import annotations

import dataclasses
from collections.abc import Callable
from typing import Any, cast

VALID_ORGANISMS: frozenset[str] = frozenset({"human", "mouse"})
VALID_DATA_SOURCES: frozenset[str] = frozenset({"sra", "gtex", "tcga"})
_VALID_GENOMIC_UNITS: frozenset[str] = frozenset({"gene", "exon"})


def _project_shard(project: str | None) -> str:
    """Return the 2-character directory shard for a project identifier.

    Project-level URLs (counts, junctions, metadata, BigWig) include a
    shard subdirectory derived from the last two characters of the project
    ID to limit directory fanout. For example, project ``SRP009615``
    produces the shard ``"15"``, yielding a path segment like
    ``.../gene_sums/15/SRP009615/...``.

    Args:
        project: Project identifier string (e.g., ``"SRP107565"``) or
            ``None``.

    Returns:
        The final two characters of ``project``. Returns the empty string
        when ``project`` is ``None`` or empty.
    """
    if not project:
        return ""
    return project[-2:]


def _sample_shard(sample: str | None, data_source: str | None) -> str:
    """Return the 2-character directory shard for a BigWig sample identifier.

    BigWig coverage file URLs include a shard subdirectory derived from the
    sample identifier. The offset within the identifier depends on the data
    source:

    * GTEx: uses positions ``[-4:-2]`` (e.g., ``"GTEX-11DXZ-0526-SM-5EQRP"``
      produces ``"EQ"``).
    * SRA / TCGA: uses the final two characters ``[-2:]`` (e.g.,
      ``"SRR387777"`` produces ``"77"``).

    Args:
        sample: Sample identifier string (e.g., ``"SRR387777"`` or
            ``"GTEX-11DXZ-0526-SM-5EQRP"``).
        data_source: Data source identifier (e.g., ``"sra"``, ``"gtex"``).

    Returns:
        A 2-character uppercase shard string, or the empty string if
        ``sample`` is ``None`` or too short.
    """
    if not sample:
        return ""
    is_gtex = isinstance(data_source, str) and "gtex" in data_source.lower()
    if is_gtex and len(sample) >= 4:
        return sample[-4:-2].upper()
    return sample[-2:].upper()


@dataclasses.dataclass(slots=True)
class _R3CommonFields:
    """Shared field schema for all resource description subclasses.

    Concrete resource description classes inherit from this dataclass to provide
    a uniform attribute surface area across all resource types. This enables
    generic code (e.g., serialization or filtering) to access common attributes
    without special-casing individual description subclasses.

    Only a subset of fields are required for a given resource type; each
    concrete subclass enforces its own required fields in `__post_init__`.

    Attributes:
        resource_type: The resource-type discriminator used by the factory.
        organism: Organism identifier (e.g., "human", "mouse").
        data_source: Source identifier (e.g., "sra", "gtex", "tcga").
        genomic_unit: Unit for feature-count matrices (e.g., "gene", "exon").
        project: Project identifier (e.g., SRP/TCGA project).
        sample: Sample/run identifier used by some resources (e.g., BigWig).
        annotation_extension: Annotation build tag (e.g., "G026").
        junction_type: Junction category selector (e.g., "ALL").
        junction_extension: Junction file format tag (e.g., "MM", "ID", "RR").
        table_name: Metadata table name (e.g., "recount_qc").
    """

    resource_type: str
    organism: str | None = None
    data_source: str | None = None
    genomic_unit: str | None = None
    project: str | None = None
    sample: str | None = None
    annotation_extension: str | None = None
    junction_type: str | None = None
    junction_extension: str | None = None
    table_name: str | None = None


class R3ResourceDescription:
    """Abstract base class and multi-factory for recount3 resource descriptors.

    Instantiate this class with a `resource_type` to obtain an instance of the
    registered concrete subclass.

    Example:
        >>> desc = R3ResourceDescription(
        ...     resource_type="annotations",
        ...     organism="human",
        ...     genomic_unit="gene",
        ...     annotation_extension="G026",
        ... )
        >>> isinstance(desc, R3Annotations)
        True

    Concrete subclasses should:
    - Inherit from both `_R3CommonFields` and `R3ResourceDescription`.
    - Validate required parameters in `__post_init__`.
    - Implement `url_path()` to return the duffel-relative path.

    Attributes:
        _TYPE_REGISTRY: Mapping from resource-type strings to concrete
            classes. Mutable to allow dynamic registration.
        _RESOURCE_TYPE: Registration string key. Injected into concrete
            subclasses by the ``register_type()`` decorator.
        resource_type: Resource-type discriminator. Declared here for
            static typing; implemented by ``_R3CommonFields``.
    """

    _TYPE_REGISTRY: dict[str, type[R3ResourceDescription]] = {}
    _RESOURCE_TYPE: str
    resource_type: str

    def __new__(
        cls,
        *args: Any,
        **kwargs: Any,
    ) -> R3ResourceDescription:
        """Constructs and returns an instance of the appropriate subclass.

        The factory accepts the resource type either as:
        - keyword argument `resource_type=...`, or
        - the first positional argument.

        Args:
            *args: Optional positional arguments; if present, `args[0]` may
                supply the resource type.
            **kwargs: Keyword arguments used to initialize the selected
                dataclass subclass.

        Returns:
            An instance of the concrete subclass registered for the selected
            resource type.

        Raises:
            KeyError: If `resource_type` is missing or empty.
            ValueError: If `resource_type` is not registered.
        """
        if cls is not R3ResourceDescription:
            return super().__new__(cls)

        resource_type = kwargs.get("resource_type")
        if resource_type is None and args:
            resource_type = args[0]

        if not resource_type:
            raise KeyError(
                "resource_type is required to construct R3ResourceDescription."
            )

        subcls = cls._TYPE_REGISTRY.get(str(resource_type))
        if subcls is None:
            raise ValueError(f"Unsupported resource_type: {resource_type!r}")

        return cast(
            "R3ResourceDescription",
            object.__new__(subcls),
        )

    @classmethod
    def register_type(
        cls,
        resource_type: str,
    ) -> Callable[[type[R3ResourceDescription]], type[R3ResourceDescription]]:
        """Registers a concrete subclass for a given `resource_type`.

        This decorator binds `resource_type` to the provided subclass and also
        sets an internal `_RESOURCE_TYPE` attribute on that subclass.

        Args:
            resource_type: Resource-type string used as the factory key.

        Returns:
            A decorator that registers the decorated subclass.

        Example:
            >>> @R3ResourceDescription.register_type("annotations")
            ... @dataclasses.dataclass(slots=True)
            ... class R3Annotations(_R3CommonFields, R3ResourceDescription):
            ...     ...
        """

        def _decorator(
            subcls: type[R3ResourceDescription],
        ) -> type[R3ResourceDescription]:
            cls._TYPE_REGISTRY[resource_type] = subcls
            subcls._RESOURCE_TYPE = resource_type  # pylint: disable=protected-access
            return subcls

        return _decorator

    def _require(self, *names: str) -> None:
        """Ensures that the given attributes are present and non-empty.

        Args:
            *names: Attribute names to validate.

        Raises:
            KeyError: If any named attribute is None or the empty string.
        """
        for name in names:
            if getattr(self, name, None) in (None, ""):
                raise KeyError(f"Missing required field: {name}")

    def _check_organism(self) -> None:
        """Validates that `organism` is one of the supported values.

        Raises:
            ValueError: If `organism` is not in `VALID_ORGANISMS`.
        """
        organism = getattr(self, "organism", None)
        if organism not in VALID_ORGANISMS:
            raise ValueError(f"Invalid organism: {organism!r}")

    def _check_data_source(self) -> None:
        """Validates that `data_source` is one of the supported values.

        Raises:
            ValueError: If `data_source` is not in `VALID_DATA_SOURCES`.
        """
        data_source = getattr(self, "data_source", None)
        if data_source not in VALID_DATA_SOURCES:
            raise ValueError(f"Invalid data_source: {data_source!r}")

    def _check_genomic_unit(self) -> None:
        """Validates that `genomic_unit` is one of the supported values.

        Raises:
            ValueError: If `genomic_unit` is not in `_VALID_GENOMIC_UNITS`.
        """
        genomic_unit = getattr(self, "genomic_unit", None)
        if genomic_unit not in _VALID_GENOMIC_UNITS:
            raise ValueError(f"Invalid genomic_unit: {genomic_unit!r}")

    def url_path(self) -> str:
        """Return the duffel-relative URL path for this resource.

        `R3ResourceDescription` is a multi-factory and abstract interface:
        calling ``R3ResourceDescription(resource_type=...)`` returns an
        instance of a concrete registered subclass (e.g. `R3Annotations`).
        Those subclasses implement `url_path()` to construct a deterministic
        path within the duffel repository layout.

        This base-class method is not expected to be called directly; it
        exists to document the interface and provide a consistent method name
        across all description types.

        Implementations must return a path without a leading slash. Callers
        should join this value onto the configured base URL (and add any
        separators as needed).

        Returns:
            A duffel-relative path string, no leading slash.

        Raises:
            NotImplementedError: Always, in the base class. Concrete subclasses
                override this method.
        """
        raise NotImplementedError


@R3ResourceDescription.register_type("annotations")
@dataclasses.dataclass(slots=True)
class R3Annotations(_R3CommonFields, R3ResourceDescription):
    """Descriptor for annotation GTF files.

    Required fields:
      - organism
      - genomic_unit
      - annotation_extension

    Duffel layout:
      {organism}/annotations/{genomic_unit}_sums/
        {organism}.{genomic_unit}_sums.{annotation_extension}.gtf.gz
    """

    def __post_init__(self) -> None:
        self._require("organism", "genomic_unit", "annotation_extension")
        self._check_organism()
        self._check_genomic_unit()

    def url_path(self) -> str:
        return (
            f"{self.organism}/annotations/{self.genomic_unit}_sums/"
            f"{self.organism}.{self.genomic_unit}_sums."
            f"{self.annotation_extension}.gtf.gz"
        )


@R3ResourceDescription.register_type("count_files_gene_or_exon")
@dataclasses.dataclass(slots=True)
class R3GeneOrExonCounts(_R3CommonFields, R3ResourceDescription):
    """Descriptor for per-project gene/exon count matrices.

    Required fields:
      - organism
      - data_source
      - genomic_unit
      - project
      - annotation_extension

    Duffel layout:
      {organism}/data_sources/{data_source}/{genomic_unit}_sums/
        {p2(project)}/{project}/
        {data_source}.{genomic_unit}_sums.{project}.{annotation_extension}.gz
    """

    def __post_init__(self) -> None:
        self._require(
            "organism",
            "data_source",
            "genomic_unit",
            "project",
            "annotation_extension",
        )
        self._check_organism()
        self._check_data_source()
        self._check_genomic_unit()

    def url_path(self) -> str:
        return (
            f"{self.organism}/data_sources/{self.data_source}/"
            f"{self.genomic_unit}_sums/{_project_shard(self.project)}/{self.project}/"
            f"{self.data_source}.{self.genomic_unit}_sums.{self.project}."
            f"{self.annotation_extension}.gz"
        )


@R3ResourceDescription.register_type("count_files_junctions")
@dataclasses.dataclass(slots=True)
class R3JunctionCounts(_R3CommonFields, R3ResourceDescription):
    """Descriptor for per-project junction count files.

    Required fields:
      - organism
      - data_source
      - project
      - junction_type
      - junction_extension

    Duffel layout:
      {organism}/data_sources/{data_source}/junctions/
        {p2(project)}/{project}/
        {data_source}.junctions.{project}.{junction_type}.{junction_extension}.gz
    """

    def __post_init__(self) -> None:
        self._require(
            "organism",
            "data_source",
            "project",
            "junction_type",
            "junction_extension",
        )
        self._check_organism()
        self._check_data_source()

    def url_path(self) -> str:
        return (
            f"{self.organism}/data_sources/{self.data_source}/junctions/"
            f"{_project_shard(self.project)}/{self.project}/"
            f"{self.data_source}.junctions.{self.project}."
            f"{self.junction_type}.{self.junction_extension}.gz"
        )


@R3ResourceDescription.register_type("metadata_files")
@dataclasses.dataclass(slots=True)
class R3ProjectMetadata(_R3CommonFields, R3ResourceDescription):
    """Descriptor for per-project metadata tables.

    Required fields:
      - organism
      - data_source
      - project
      - table_name

    Duffel layout:
      {organism}/data_sources/{data_source}/metadata/
        {p2(project)}/{project}/
        {data_source}.{table_name}.{project}.MD.gz
    """

    def __post_init__(self) -> None:
        self._require("organism", "data_source", "project", "table_name")
        self._check_organism()
        self._check_data_source()

    def url_path(self) -> str:
        return (
            f"{self.organism}/data_sources/{self.data_source}/metadata/"
            f"{_project_shard(self.project)}/{self.project}/"
            f"{self.data_source}.{self.table_name}.{self.project}.MD.gz"
        )


@R3ResourceDescription.register_type("bigwig_files")
@dataclasses.dataclass(slots=True)
class R3BigWig(_R3CommonFields, R3ResourceDescription):
    """Descriptor for per-sample BigWig coverage files.

    Required fields:
      - organism
      - data_source
      - project
      - sample

    Duffel layout:
      {organism}/data_sources/{data_source}/base_sums/
        {p2(project)}/{project}/{shard(sample, data_source)}/
        {data_source}.base_sums.{project}_{sample}.ALL.bw

    The sample shard subdirectory uses a different offset for GTEx samples
    compared to SRA/TCGA samples; see :func:`_sample_shard`.
    """

    def __post_init__(self) -> None:
        self._require("organism", "data_source", "project", "sample")
        self._check_organism()
        self._check_data_source()

    def url_path(self) -> str:
        shard = _sample_shard(self.sample, self.data_source)
        return (
            f"{self.organism}/data_sources/{self.data_source}/base_sums/"
            f"{_project_shard(self.project)}/{self.project}/{shard}/"
            f"{self.data_source}.base_sums.{self.project}_{self.sample}.ALL.bw"
        )


@R3ResourceDescription.register_type("data_sources")
@dataclasses.dataclass(slots=True)
class R3DataSources(_R3CommonFields, R3ResourceDescription):
    """Descriptor for the organism-level data-source index (homes_index).

    Required fields:
      - organism

    Duffel layout:
      {organism}/homes_index
    """

    def __post_init__(self) -> None:
        self._require("organism")
        self._check_organism()

    def url_path(self) -> str:
        return f"{self.organism}/homes_index"


@R3ResourceDescription.register_type("data_source_metadata")
@dataclasses.dataclass(slots=True)
class R3DataSourceMetadata(_R3CommonFields, R3ResourceDescription):
    """Descriptor for source-level metadata listings.

    Required fields:
      - organism
      - data_source

    Duffel layout:
      {organism}/data_sources/{data_source}/metadata/
        {data_source}.recount_project.MD.gz
    """

    def __post_init__(self) -> None:
        self._require("organism", "data_source")
        self._check_organism()
        self._check_data_source()

    def url_path(self) -> str:
        return (
            f"{self.organism}/data_sources/{self.data_source}/metadata/"
            f"{self.data_source}.recount_project.MD.gz"
        )
