import json
import os

BASE_DIR = os.path.dirname(os.path.abspath(__file__))


def load_config():
    cfg_path = os.path.join(BASE_DIR, "config.json")
    # if not os.path.exists(cfg_path):
    #     raise FileNotFoundError(f"config.json file not found at {cfg_path}. Please create this file with the required configuration.")
    with open(cfg_path, "r") as f:
        return json.load(f)


def save_config(config, config_path=None):
    cfg_path = config_path if config_path is not None else os.path.join(BASE_DIR, "config.json")
    # Ensure target directory exists when writing to a custom location
    os.makedirs(os.path.dirname(cfg_path), exist_ok=True)
    with open(cfg_path, "w") as f:
        json.dump(config, f, indent=4)


def compute_decennial_year(main_year):
    try:
        return 2020 if int(main_year) >= 2020 else 2010
    except Exception:
        return 2010


def update_config_values(config,
                         census_api_key=None,
                         main_year=None,
                         geos=None,
                         commute_states=None,
                         use_pums=None,
                         path=None,
                         julia_env_path=None,
                         LODES_annual_income_boundary=None,
                         income_associativity_coefficient=None,
                         school_associativity_coefficient=None,
                         inst_res_per_worker=None,
                         noninst_res_per_worker=None,
                         min_gq_workers=None,
                         min_gq_residents=None,
                         n_closest_schools=None,
                         p_closest_school=None,
                         CO_crit_val=None,
                         CO_cooldown=None,
                         CO_maxgens=None,
                         workplace_K=None,
                         school_K=None,
                         gq_K=None,
                         netw_K=None,
                         netw_B=None):
    if census_api_key is not None:
        config["census_api_key"] = census_api_key
    if main_year is not None:
        config["main_year"] = main_year
        config["decennial_year"] = compute_decennial_year(main_year)
    if geos is not None:
        config["geos"] = geos
    if commute_states is not None:
        config["commute_states"] = commute_states
    if use_pums is not None:
        config["use_pums"] = use_pums
    if path is not None:
        config["path"] = path
    if julia_env_path is not None:
        config["julia_env_path"] = julia_env_path
    if LODES_annual_income_boundary is not None:
        config["LODES_annual_income_boundary"] = LODES_annual_income_boundary
    if income_associativity_coefficient is not None:
        config["income_associativity_coefficient"] = income_associativity_coefficient
    if school_associativity_coefficient is not None:
        config["school_associativity_coefficient"] = school_associativity_coefficient
    if inst_res_per_worker is not None:
        config["inst_res_per_worker"] = inst_res_per_worker
    if noninst_res_per_worker is not None:
        config["noninst_res_per_worker"] = noninst_res_per_worker
    if min_gq_workers is not None:
        config["min_gq_workers"] = min_gq_workers
    if min_gq_residents is not None:
        config["min_gq_residents"] = min_gq_residents
    if n_closest_schools is not None:
        config["n_closest_schools"] = n_closest_schools
    if p_closest_school is not None:
        config["p_closest_school"] = p_closest_school
    if CO_crit_val is not None:
        config["CO_crit_val"] = CO_crit_val
    if CO_cooldown is not None:
        config["CO_cooldown"] = CO_cooldown
    if CO_maxgens is not None:
        config["CO_maxgens"] = CO_maxgens
    if workplace_K is not None:
        config["workplace_K"] = workplace_K
    if school_K is not None:
        config["school_K"] = school_K
    if gq_K is not None:
        config["gq_K"] = gq_K
    if netw_K is not None:
        config["netw_K"] = netw_K
    if netw_B is not None:
        config["netw_B"] = netw_B
    return config


class WriteConfig:
    def __init__(self,
                 census_api_key=None,
                 main_year=None,
                 geos=None,
                 commute_states=None,
                 use_pums=None,
                 path=None,
                 julia_env_path=None,
                 config_dict=None,
                 base_dir=None,
                 LODES_annual_income_boundary=None,
                 income_associativity_coefficient=None,
                 school_associativity_coefficient=None,
                 inst_res_per_worker=None,
                 noninst_res_per_worker=None,
                 min_gq_workers=None,
                 min_gq_residents=None,
                 n_closest_schools=None,
                 p_closest_school=None,
                 CO_crit_val=None,
                 CO_cooldown=None,
                 CO_maxgens=None,
                 workplace_K=None,
                 school_K=None,
                 gq_K=None,
                 netw_K=None,
                 netw_B=None):
        self.base_dir = base_dir if base_dir is not None else BASE_DIR
        # Load base template from the package directory unless a dict is provided
        self.template_config_path = os.path.join(self.base_dir, "config.json")
        if path is not None:
            # If path is a directory, append 'config.json' to it
            if os.path.isdir(path):
                self.path = os.path.join(path, "config.json")
            else:
                self.path = path
        else:
            self.path = self.template_config_path
        self.config = config_dict if config_dict is not None else load_config()
        self.overrides = {
            "census_api_key": census_api_key,
            "main_year": main_year,
            "geos": geos,
            "commute_states": commute_states,
            "use_pums": use_pums,
            "path": path,
            "julia_env_path": julia_env_path,
            "LODES_annual_income_boundary": LODES_annual_income_boundary,
            "income_associativity_coefficient": income_associativity_coefficient,
            "school_associativity_coefficient": school_associativity_coefficient,
            "inst_res_per_worker": inst_res_per_worker,
            "noninst_res_per_worker": noninst_res_per_worker,
            "min_gq_workers": min_gq_workers,
            "min_gq_residents": min_gq_residents,
            "n_closest_schools": n_closest_schools,
            "p_closest_school": p_closest_school,
            "CO_crit_val": CO_crit_val,
            "CO_cooldown": CO_cooldown,
            "CO_maxgens": CO_maxgens,
            "workplace_K": workplace_K,
            "school_K": school_K,
            "gq_K": gq_K,
            "netw_K": netw_K,
            "netw_B": netw_B,
        }
        
        self.run_all()

    def run_all(self):
        update_config_values(
            self.config,
            census_api_key=self.overrides["census_api_key"],
            main_year=self.overrides["main_year"],
            geos=self.overrides["geos"],
            commute_states=self.overrides["commute_states"],
            use_pums=self.overrides["use_pums"],
            path=self.overrides["path"],
            julia_env_path=self.overrides["julia_env_path"],
            LODES_annual_income_boundary=self.overrides["LODES_annual_income_boundary"],
            income_associativity_coefficient=self.overrides["income_associativity_coefficient"],
            school_associativity_coefficient=self.overrides["school_associativity_coefficient"],
            inst_res_per_worker=self.overrides["inst_res_per_worker"],
            noninst_res_per_worker=self.overrides["noninst_res_per_worker"],
            min_gq_workers=self.overrides["min_gq_workers"],
            min_gq_residents=self.overrides["min_gq_residents"],
            n_closest_schools=self.overrides["n_closest_schools"],
            p_closest_school=self.overrides["p_closest_school"],
            CO_crit_val=self.overrides["CO_crit_val"],
            CO_cooldown=self.overrides["CO_cooldown"],
            CO_maxgens=self.overrides["CO_maxgens"],
            workplace_K=self.overrides["workplace_K"],
            school_K=self.overrides["school_K"],
            gq_K=self.overrides["gq_K"],
            netw_K=self.overrides["netw_K"],
            netw_B=self.overrides["netw_B"],
        )
        # Save to user-specified path
        save_config(self.config, self.path)
        # Also save to package directory
        save_config(self.config, self.template_config_path)
        print("Updated config.json with new parameters")

    def get_pars(self):
        with open(self.path, "r") as f:
            cfg = json.load(f)
        print(json.dumps(cfg, indent=2))


def main():
    runner = WriteConfig()
    runner.run_all()


if __name__ == "__main__":
    main()