import itertools
from os import PathLike
from pathlib import Path
from typing import cast

import polars as pl

from matchescu.extraction import Traits
from matchescu.extraction.csv import CsvRecordExtraction
from matchescu.matching.config import MagellanBenchmarkDataConfig, TraitConfig
from matchescu.matching.evaluation.data.benchmark import BenchmarkDataBuilder
from matchescu.matching.evaluation.data.benchmark._base import BenchmarkData
from matchescu.matching.evaluation.data.benchmark._config_adapters import get_traits
from matchescu.matching.evaluation.data.splits._split import Split
from matchescu.matching.evaluation.ground_truth import EquivalenceClassPartitioner
from matchescu.reference_store.comparison_space import InMemoryComparisonSpace
from matchescu.typing import EntityReferenceIdentifier as RefId


class MagellanBenchmarkData(BenchmarkData):
    def __init__(self, folder_path: str | PathLike) -> None:
        super().__init__()
        self.__dataset_dir = Path(folder_path)
        if not self.__dataset_dir.exists():
            raise FileNotFoundError(self.__dataset_dir)
        if not self.__dataset_dir.is_dir():
            raise NotADirectoryError(self.__dataset_dir)
        self.__left_table_path = self.__dataset_dir / "tableA.csv"
        self.__left_source = self.__left_table_path.stem
        self.__right_table_path = self.__dataset_dir / "tableB.csv"
        self.__right_source = self.__right_table_path.stem
        self.__split_paths = [
            self.__dataset_dir / f"{split}.csv" for split in self._SPLIT_NAMES
        ]
        self._check_files(
            [
                self.__left_table_path,
                self.__right_table_path,
                *self.__split_paths,
            ]
        )

    def _load_csv_table(self, path: Path, traits: Traits) -> str:
        extract = CsvRecordExtraction(path, traits)
        for ref in extract():
            self._id_table.put(ref)
        return extract.data_source.name

    def load_left(self, traits: Traits) -> "MagellanBenchmarkData":
        self.__left_source = self._load_csv_table(self.__left_table_path, traits)
        return self

    def load_right(self, traits: Traits) -> "MagellanBenchmarkData":
        self.__right_source = self._load_csv_table(self.__right_table_path, traits)
        return self

    def __load_split(self, path: Path) -> Split:
        rows = list(
            itertools.starmap(
                lambda x, y, is_match: (
                    RefId(x, self.__left_source),
                    RefId(y, self.__right_source),
                    is_match,
                ),
                pl.read_csv(path, ignore_errors=True).iter_rows(named=False),
            )
        )
        cs = InMemoryComparisonSpace()
        ids = {}
        for left, right, _ in rows:
            cs.put(left, right)
            ids[left] = None
            ids[right] = None

        matcher_gt = {(left, right): label for left, right, label in rows if label > 0}
        self._match_gt.update(matcher_gt)
        ecp = EquivalenceClassPartitioner(ids)
        clusters = {
            cluster_no: set(cluster)
            for cluster_no, cluster in enumerate(ecp(matcher_gt), 1)
        }
        return Split(cs, matcher_gt, clusters)

    def load_splits(self) -> "MagellanBenchmarkData":
        if not self.__left_source or not self.__right_source:
            raise RuntimeError("load left and right tables before splitting")
        del self._match_gt
        del self._cluster_gt
        self._match_gt = {}
        self._cluster_gt = {}
        self._splits = {
            f"{name}_split": self.__load_split(path)
            for name, path in zip(self._SPLIT_NAMES, self.__split_paths)
        }
        ecp = EquivalenceClassPartitioner(self._id_table.ids())
        self._cluster_gt = {
            ref_id: cluster_no
            for cluster_no, cluster in enumerate(ecp(self._match_gt), start=1)
            for ref_id in cluster
        }
        return self

    @property
    def name(self) -> str:
        return self.__dataset_dir.stem

    @property
    def left_source(self) -> str:
        return self.__left_source

    @property
    def right_source(self) -> str:
        return self.__right_source

    def __getattr__(self, item: str):
        if item in self._splits:
            return self._splits[item]
        raise AttributeError(f"{item} split not found")

    def __dir__(self):
        return sorted([*super().__dir__(), *self._splits.keys()])

    def all_data(self) -> Split:
        return self._global_split


class MagellanTraits:
    __TRAIT_DICT = {
        "ABT-BUY": Traits().string(["name", "description"]).currency(["price"]),
        "AMAZON-GOOGLE": Traits().string(["title", "manufacturer"]).currency(["price"]),
        "BEER": Traits().string(["Beer_Name", "Brew_Factory_Name", "Style"]),
        "DBLP-ACM": Traits().string(["title", "authors", "venue"]).int(["year"]),
        "DBLP-SCHOLAR": Traits().string(["title", "authors", "venue", "year"]),
        "FODORS-ZAGAT": Traits()
        .string(["name", "addr", "city", "phone", "type"])
        .int(["class"]),
        "ITUNES-AMAZON": Traits().string(
            [
                "Song_Name",
                "Artist_Name",
                "Album_Name",
                "Genre",
                "Price",
                "CopyRight",
                "Time",
                "Released",
            ]
        ),
        "WALMART-AMAZON": Traits()
        .string(["title", "category", "brand", "modelno"])
        .currency(["price"]),
    }

    def __getitem__(self, item: str) -> Traits:
        assert isinstance(item, str)
        key = item.upper()
        assert key in self.__TRAIT_DICT
        return self.__TRAIT_DICT[key]


class MagellanBenchmarkDataBuilder(BenchmarkDataBuilder[MagellanBenchmarkData]):
    """Initialize fully loaded ``MagellanBenchmarkData`` from config."""

    _MAGELLAN_TRAITS = MagellanTraits()

    def __init__(
        self, params: MagellanBenchmarkDataConfig, data_dir: Path | None = None
    ) -> None:
        super().__init__(params, data_dir)
        self._params = cast(MagellanBenchmarkDataConfig, self._params)

    def _create_instance(self) -> MagellanBenchmarkData:
        return MagellanBenchmarkData(self._data_dir)

    @classmethod
    def _resolve_traits(cls, trait_configs: str | list[TraitConfig]):
        """A string is looked up in the well-known MagellanTraits dict;
        a list of TraitConfig is built manually."""
        if isinstance(trait_configs, str):
            return cls._MAGELLAN_TRAITS[trait_configs]
        return get_traits(trait_configs)

    def load_data(self) -> BenchmarkDataBuilder[MagellanBenchmarkData]:
        left_traits = self._resolve_traits(self._params.left_traits)
        right_traits = (
            self._resolve_traits(self._params.right_traits)
            if self._params.right_traits is not None
            else left_traits
        )
        self._instance.load_left(left_traits).load_right(right_traits)
        return self

    def load_splits(self) -> BenchmarkDataBuilder[MagellanBenchmarkData]:
        self._instance.load_splits()
        return self
