Analyzer interface

class AssetsReader(ABC):
    @abstractmethod
    def table(self, output_id: str) -> "TableReader":
        """
        Gets the table reader for the specified output.
        """
        pass

class BaseDerivedModuleContext(ABC, BaseModel):
    """
    Common interface for secondary analyzers runtime contexts.
    """

    temp_dir: str
    """
  Gets the temporary directory that the module can freely write content to
  during its lifetime. This directory will not persist between runs.
  """

    @property
    @abstractmethod
    def base_params(self) -> dict[str, ParamValue]:
        """
        Gets the primary analysis parameters.
        """
        pass

    @property
    @abstractmethod
    def base(self) -> "AssetsReader":
        """
        Gets the base primary analyzer's context, which lets you inspect and load its
        outputs.
        """
        pass

    @abstractmethod
    def dependency(
        self, secondary_interface: SecondaryAnalyzerInterface
    ) -> "AssetsReader":
        """
        Gets the context of a secondary analyzer the current module depends on, which
        lets you inspect and load its outputs.
        """
        pass

class InputTableReader(TableReader):
    @abstractmethod
    def preprocess[PolarsDataFrameLike](
        self, df: PolarsDataFrameLike
    ) -> PolarsDataFrameLike:
        """
        Given the manually loaded user input dataframe, apply column mapping and
        semantic transformations to give the input dataframe that the analyzer
        expects.
        """
        pass

class NullProgressReporter:
    """No-op progress reporter for when progress reporting is disabled."""

    def __init__(self, title: str):
        self.title = title

    def update(self, value: float):
        pass

    def finish(self, done_text: str = "Done!"):
        pass

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        pass

class PrimaryAnalyzerContext(ABC, BaseModel):
    temp_dir: str
    """
  Gets the temporary directory that the module can freely write content to
  during its lifetime. This directory will not persist between runs.
  """

    @abstractmethod
    def input(self) -> "InputTableReader":
        """
        Gets the input reader context.

        **Note that this is in function form** even though one input is expected,
        in anticipation that we may want to support multiple inputs in the future.
        """
        pass

    @property
    @abstractmethod
    def params(self) -> dict[str, ParamValue]:
        """
        Gets the analysis parameters.
        """
        pass

    @abstractmethod
    def output(self, output_id: str) -> "TableWriter":
        """
        Gets the output writer context for the specified output ID.
        """
        pass

class ProgressReporterProtocol(Protocol):
    """Protocol for progress reporting during analysis steps."""

    def update(self, value: float) -> None:
        """Update progress value (0.0 to 1.0)."""
        ...

    def finish(self, done_text: str = "Done!") -> None:
        """Mark the step as complete."""
        ...

    def __enter__(self) -> "ProgressReporterProtocol": ...

    def __exit__(self, *args) -> None: ...

class SecondaryAnalyzerContext(BaseDerivedModuleContext):
    @abstractmethod
    def output(self, output_id: str) -> "TableWriter":
        """
        Gets the output writer context
        """
        pass

class TableReader(ABC):
    @property
    @abstractmethod
    def parquet_path(self) -> str:
        """
        Gets the path to the table's parquet file. The module should expect a parquet
        file here.
        """
        pass

class TableWriter(ABC):
    @property
    @abstractmethod
    def parquet_path(self) -> str:
        """
        Gets the path to the table's parquet file. The module should write a parquet
        file to it.
        """
        pass

class AnalyzerDeclaration(AnalyzerInterface):
    entry_point: Callable[[PrimaryAnalyzerContext], None]
    default_params: Callable[[PrimaryAnalyzerContext], dict[str, ParamValue]]
    is_distributed: bool

    def __init__(
        self,
        interface: AnalyzerInterface,
        main: Callable,
        *,
        is_distributed: bool = False,
        default_params: Callable[[PrimaryAnalyzerContext], dict[str, ParamValue]] = (
            lambda _: dict()
        ),
    ):
        """Creates a primary analyzer declaration

        Args:
          interface (AnalyzerInterface): The metadata interface for the primary analyzer.

          main (Callable):
            The entry point function for the primary analyzer. This function should
            take a single argument of type `PrimaryAnalyzerContext` and should ensure
            that the outputs specified in the interface are generated.

          is_distributed (bool):
            Set this explicitly to `True` once the analyzer is ready to be shipped
            to end users; it will make the analyzer available in the distributed
            executable.
        """
        super().__init__(
            **interface.model_dump(),
            entry_point=main,
            default_params=default_params,
            is_distributed=is_distributed,
        )

class SecondaryAnalyzerDeclaration(SecondaryAnalyzerInterface):
    entry_point: Callable[["SecondaryAnalyzerContext"], None]

    def __init__(self, interface: SecondaryAnalyzerInterface, main: Callable):
        """Creates a secondary analyzer declaration

        Args:
          interface (SecondaryAnalyzerInterface): The metadata interface for the secondary analyzer.

          main (Callable):
            The entry point function for the secondary analyzer. This function should
            take a single argument of type `SecondaryAnalyzerContext` and should ensure
            that the outputs specified in the interface are generated.
        """
        super().__init__(**interface.model_dump(), entry_point=main)

class AnalyzerInterface(BaseAnalyzerInterface):
    input: AnalyzerInput
    """
  Specifies the input data schema for the analyzer.
  """

    params: list[AnalyzerParam] = []
    """
  A list of parameters that the analyzer accepts.
  """

    outputs: list["AnalyzerOutput"]
    """
  Specifies the output data schema for the analyzer.
  """

    kind: Literal["primary"] = "primary"

class AnalyzerOutput(BaseModel):
    id: str
    """
  Uniquely identifies the output data schema for the analyzer. The analyzer
  must include this key in the output dictionary.
  """

    name: str
    """The human-friendly for the output."""

    description: Optional[str] = None

    columns: list["OutputColumn"]

    internal: bool = False

    def get_column_by_name(self, name: str):
        for column in self.columns:
            if column.name == name:
                return column
        return None

    def transform_output(self, output_df: pl.LazyFrame | pl.DataFrame):
        output_columns = output_df.lazy().collect_schema().names()
        return output_df.select(
            [
                pl.col(col_name).alias(
                    output_spec.human_readable_name_or_fallback()
                    if output_spec
                    else col_name
                )
                for col_name in output_columns
                if (output_spec := self.get_column_by_name(col_name)) or True
            ]
        )

class AnalyzerParam(BaseModel):
    id: str
    """
    The name of the parameter. This becomes the key in the parameters dictionary
    that is passed to the analyzer.
    """

    human_readable_name: Optional[str] = None
    """
    The human-friendly name for the parameter. This is used in the UI to
    represent the parameter.
    """

    description: Optional[str] = None
    """
    A short description of the parameter. This is used in the UI to represent
    the parameter.
    """

    type: ParamType
    """
    The type of the parameter. This is used for validation and for customizing
    the UX for parameter input.
    """

    default: Optional[ParamValue] = None
    """
    Optional: define a static default value for this parameter. A parameter
    without a default will need to be chosen explicitly by the user.
    """

    backfill_value: Optional[ParamValue] = None
    """
    Recommended if this is a parameter that is newly introduced in a previously
    released analyzer. The backfill is show what this parameter was before it
    became customizable.
    """

    @property
    def print_name(self):
        return self.human_readable_name or self.id

class BaseAnalyzerInterface(BaseModel):
    id: str
    """
  The static ID for the analyzer that, with the version, uniquely identifies the
  analyzer and will be stored as metadata as part of the output data.
  """

    version: str
    """
  The version ID for the analyzer. In future, we may choose to support output
  migration between versions of the same analyzer.
  """

    name: str
    """
  The short human-readable name of the analyzer.
  """

    short_description: str
    """
  A short, one-liner description of what the analyzer does.
  """

    long_description: Optional[str] = None
    """
  A longer description of what the analyzer does that will be shown separately.
  """

class DerivedAnalyzerInterface(BaseAnalyzerInterface):
    base_analyzer: AnalyzerInterface
    """
  The base analyzer that this secondary analyzer extends. This is always a primary
  analyzer. If your module depends on other secondary analyzers (which must have
  the same base analyzer), you can specify them in the `depends_on` field.
  """

    depends_on: list["SecondaryAnalyzerInterface"] = []
    """
  A dictionary of secondary analyzers that must be run before the current analyzer
  secondary analyzer is run. These secondary analyzers must have the same
  primary base.
  """

class InputColumn(Column):
    name_hints: list[str] = []
    """
  Specifies a list of space-separated words that are likely to be found in the
  column name of the user-provided data. This is used to help the user map the
  input columns to the expected columns.

  Any individual hint matching is sufficient for a match to be called. The hint
  in turn is matched if every word matches some part of the column name.
  """

class SecondaryAnalyzerInterface(DerivedAnalyzerInterface):
    outputs: list[AnalyzerOutput]
    """
  Specifies the output data schema for the analyzer.
  """

    kind: Literal["secondary"] = "secondary"

class IntegerParam(BaseModel):
    """
    Represents an integer value

    The corresponding value will be of type `int`.
    """

    type: Literal["integer"] = "integer"
    min: int
    max: int

class TimeBinningParam(BaseModel):
    """
    Represents a time bin.

    The corresponding value will be of type `TimeBinningValue`.
    """

    type: Literal["time_binning"] = "time_binning"

class TimeBinningValue(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    unit: TimeBinningUnit
    amount: int

    def to_polars_truncate_spec(self) -> str:
        """
        Converts the value to a string that can be used in Polars truncate spec.
        See https://docs.pola.rs/api/python/stable/reference/expressions/api/polars.Expr.dt.truncate.html
        """
        amount = self.amount
        unit = self.unit
        if unit == "year":
            return f"{amount}y"
        if unit == "month":
            return f"{amount}mo"
        if unit == "week":
            return f"{amount}w"
        if unit == "day":
            return f"{amount}d"
        if unit == "hour":
            return f"{amount}h"
        if unit == "minute":
            return f"{amount}m"
        if unit == "second":
            return f"{amount}s"

        raise ValueError("Invalid time binning value")

    def to_human_readable_text(self) -> str:
        amount = self.amount
        unit = self.unit

        if unit == "year":
            return f"{amount} year{'s' if amount > 1 else ''}"
        if unit == "month":
            return f"{amount} month{'s' if amount > 1 else ''}"
        if unit == "week":
            return f"{amount} week{'s' if amount > 1 else ''}"
        if unit == "day":
            return f"{amount} day{'s' if amount > 1 else ''}"
        if unit == "hour":
            return f"{amount} hour{'s' if amount > 1 else ''}"
        if unit == "minute":
            return f"{amount} minute{'s' if amount > 1 else ''}"
        if unit == "second":
            return f"{amount} second{'s' if amount > 1 else ''}"

        raise ValueError("Invalid time binning value")