from __future__ import annotations

from typing import Any

import pandas as pd

from sdgx.data_models.inspectors.base import Inspector
from sdgx.data_models.inspectors.extension import hookimpl


class IDInspector(Inspector):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.ID_columns: set[str] = set()

    def fit(self, raw_data: pd.DataFrame, *args, **kwargs):
        """Fit the inspector.

        Gets the list of discrete columns from the raw data.

        Args:
            raw_data (pd.DataFrame): Raw data
        """

        df_length = len(raw_data)
        candidate_columns = set(raw_data.select_dtypes(include=["object", "int64"]).columns)

        for each_col_name in candidate_columns:
            target_col = raw_data[each_col_name]
            col_set_length = len(set(target_col))
            if col_set_length == df_length:
                self.ID_columns.add(each_col_name)

        self.ready = True

    def inspect(self, *args, **kwargs) -> dict[str, Any]:
        """Inspect raw data and generate metadata."""

        return {"id_columns": list(self.ID_columns)}


@hookimpl
def register(manager):
    manager.register("IDInspector", IDInspector)
