from hashlib import sha256
import inspect
import re
from typing import List

from datalabs.packaged_modules.csv import csv
from datalabs.packaged_modules.json import json
from datalabs.packaged_modules.pandas import pandas
from datalabs.packaged_modules.parquet import parquet
from datalabs.packaged_modules.text import text


def hash_python_lines(lines: List[str]) -> str:
    filtered_lines = []
    for line in lines:
        line = re.sub(r"#.*", "", line)  # remove comments
        if line:
            filtered_lines.append(line)
    full_str = "\n".join(filtered_lines)

    # Make a hash from all this code
    full_bytes = full_str.encode("utf-8")
    return sha256(full_bytes).hexdigest()


# get importable module names and hash for caching
_PACKAGED_DATASETS_MODULES = {
    "csv": (csv.__name__, hash_python_lines(inspect.getsource(csv).splitlines())),
    "json": (json.__name__, hash_python_lines(inspect.getsource(json).splitlines())),
    "pandas": (
        pandas.__name__,
        hash_python_lines(inspect.getsource(pandas).splitlines()),
    ),
    "parquet": (
        parquet.__name__,
        hash_python_lines(inspect.getsource(parquet).splitlines()),
    ),
    "text": (text.__name__, hash_python_lines(inspect.getsource(text).splitlines())),
}

_EXTENSION_TO_MODULE = {
    "csv": "csv",
    "tsv": "csv",
    "json": "json",
    "jsonl": "json",
    "parquet": "parquet",
    "txt": "text",
}
