# MIT License

# Copyright (c) 2025 Ashlin Iser, Karlsruhe Institute of Technology (KIT)

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.

import pandas as pd
import os
import glob
import warnings
import polars as pl

from gbd_core.contexts import suffixes, identify, get_context_by_suffix
from gbd_core.api import GBD, GBDException
from gbd_core.util import eprint, confirm
from gbd_init.initializer import Initializer, InitializerException

try:
    from gbdc import (
        extract_base_features,
        base_feature_names,
        extract_gate_features,
        gate_feature_names,
        isohash,
        wcnfisohash,
        wcnf_base_feature_names,
        extract_wcnf_base_features,
        opb_base_feature_names,
        extract_opb_base_features,
        checksani,
        checksani_feature_names,
    )
except ImportError:

    def extract_base_features(path, tlim, mlim):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def base_feature_names():
        return []

    def extract_gate_features(path, tlim, mlim):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def gate_feature_names():
        return []

    def isohash(path):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def extract_wcnf_base_features(path, tlim, mlim):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def wcnf_base_feature_names():
        return []

    def extract_opb_base_features(path, tlim, mlim):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def opb_base_feature_names():
        return []

    def checksani(path, tlim, mlim):
        raise ModuleNotFoundError("gbdc not found", name="gbdc")

    def checksani_feature_names():
        return []


## GBDHash
def compute_hash(hash, path, limits):
    eprint("Hashing {}".format(path))
    hash = identify(path)
    return [("local", hash, path), ("filename", hash, os.path.basename(path))]


## ISOHash
def compute_isohash(hash, path, limits):
    eprint("Computing ISOHash for {}".format(path))
    context = get_context_by_suffix(path)
    if context == "wcnf":
        ihash = wcnfisohash(path)
    else:
        ihash = isohash(path)
    return [("isohash", hash, ihash)]


## Base Features
def compute_base_features(hash, path, limits, tp=None):
    eprint("Extracting base features from {} {}".format(hash, path))
    rec = extract_base_features(path, limits["tlim"], limits["mlim"])
    return [(key, hash, int(value) if isinstance(value, float) and value.is_integer() else value) for key, value in rec.items()]


## Gate Features
def compute_gate_features(hash, path, limits, tp=None):
    eprint("Extracting gate features from {} {}".format(hash, path))
    rec = extract_gate_features(path, limits["tlim"], limits["mlim"])
    return [(key, hash, int(value) if isinstance(value, float) and value.is_integer() else value) for key, value in rec.items()]


## WCNF Base Features
def compute_wcnf_base_features(hash, path, limits, tp=None):
    eprint("Extracting WCNF base features from {} {}".format(hash, path))
    rec = extract_wcnf_base_features(path, limits["tlim"], limits["mlim"])
    return [(key, hash, int(value) if isinstance(value, float) and value.is_integer() else value) for key, value in rec.items()]


## OPB Base Features
def compute_opb_base_features(hash, path, limits, tp=None):
    eprint("Extracting OPB base features from {} {}".format(hash, path))
    rec = extract_opb_base_features(path, limits["tlim"], limits["mlim"])
    return [(key, hash, int(value) if isinstance(value, float) and value.is_integer() else value) for key, value in rec.items()]


## SANI Features
def compute_sani_features(hash, path, limits, tp=None):
    eprint("Extracting SANI features from {} {}".format(hash, path))
    rec = checksani(path, limits["tlim"], limits["mlim"])
    return [(key, hash, int(value) if isinstance(value, float) and value.is_integer() else value) for key, value in rec.items()]


generic_extractors = {
    "base": {
        "description": "Extract base features from CNF files. ",
        "contexts": ["cnf", "sancnf"],
        "features": [(name, "empty") for name in base_feature_names()],
        "compute": compute_base_features,
    },
    "checksani": {
        "description": "Extract sanitise status from CNF files. ",
        "contexts": ["cnf", "sancnf"],
        "features": [(name, "empty") for name in checksani_feature_names()],
        "compute": compute_sani_features,
    },
    "gate": {
        "description": "Extract gate features from CNF files. ",
        "contexts": ["cnf", "sancnf"],
        "features": [(name, "empty") for name in gate_feature_names()],
        "compute": compute_gate_features,
    },
    "isohash": {
        "description": "Compute ISOHash for CNF or WCNF files. ",
        "contexts": ["cnf", "wcnf", "sancnf"],
        "features": [("isohash", "empty")],
        "compute": compute_isohash,
    },
    "wcnfbase": {
        "description": "Extract base features from WCNF files. ",
        "contexts": ["wcnf"],
        "features": [(name, "empty") for name in wcnf_base_feature_names()],
        "compute": compute_wcnf_base_features,
    },
    "opbbase": {
        "description": "Extract base features from OPB files. ",
        "contexts": ["opb"],
        "features": [(name, "empty") for name in opb_base_feature_names()],
        "compute": compute_opb_base_features,
    },
}


def init_features_generic(key: str, api: GBD, rlimits, df: pl.DataFrame, target_db):
    einfo = generic_extractors[key]
    context = api.database.dcontext(target_db)
    if not context in einfo["contexts"]:
        raise InitializerException("Target database context must be in {}".format(einfo["contexts"]))
    extractor = Initializer(api, rlimits, target_db, einfo["features"], einfo["compute"])
    extractor.create_features()
    extractor.run(df)


def init_local(api: GBD, rlimits, root, target_db):
    context = api.database.dcontext(target_db)

    features = [("local", None), ("filename", None)]
    extractor = Initializer(api, rlimits, target_db, features, compute_hash)
    extractor.create_features()

    # Cleanup stale entries
    df: pl.DataFrame = api.query(group_by=context + ":local", collapse=None)
     
    def path_exists(p):
        return p is not None and os.path.exists(p)
    
    missing = df.with_columns(
        exists=pl.col("local").map_elements(
            path_exists,
            return_dtype=pl.Boolean
        )
    ).filter(~pl.col("exists")).select("local")
    
    if len(missing) and api.verbose:
        for path in missing["local"].to_list():
            eprint(path)
    if len(missing) and confirm("{} files not found. Remove stale entries from local table?".format(len(missing))):
        api.reset_values("local", values=missing["local"].to_list())

    # Create df with paths not yet in local table
    paths = [path for suffix in suffixes(context) for path in glob.iglob(root + "/**/*" + suffix, recursive=True)]
    df2 = pl.DataFrame([(None, path) for path in paths if path not in df["local"].to_list()], schema=["hash", "local"], orient="row")

    extractor.run(df2)
