import os.path as osp
import yaml
import warnings
from copy import deepcopy
from easydict import EasyDict as edict

    
def load(
    experiment_file,
    *,
    module_files=(),
    base_dir=".",
    return_context:bool=False):
    ctx = {}

    mods = {}
    for pth in module_files:
        if pth.endswith(".yml"):
            mod_name = osp.basename(pth)[:-4]
            mod_pth = osp.join(base_dir, pth)
            with open(mod_pth, "r") as f:
                mods[mod_name] = _load_mod(f, mods, ctx, mod_name, deepcopy_=False)
        else:
            raise ValueError(f"Unsupported file type: {pth}")
    
    exp_pth = osp.join(base_dir, experiment_file)
    if exp_pth.endswith(".yml"):
        with open(exp_pth, "r") as f:
            exp_name = osp.basename(pth)[:-4]
            _exp = _load_mod(f, mods, ctx, exp_name, deepcopy_=True)
    else:
        raise ValueError(f"Unsupported file type: {exp_pth}")
    
    if return_context:
        return _exp, mods, ctx
    else:
        return edict(_exp.official)


def _load_mod(f, mods, ctx, ns, deepcopy_):
    newmod = yaml.safe_load(f)
    newmod = AliasDict(newmod, mods, ctx, ns, deepcopy_)
    return newmod


def _resolve_dollar(dct:dict, mods, ctx, deepcopy_):
    for key, value in dct.items():
        if isinstance(value, str) and value.startswith("$"):
            ref_key_end = value.find("(")
            if ref_key_end < 0:
                ref_key_end = len(value)
            ref_key = value[1:ref_key_end]
            
            if ref_key_end != len(value):
                ref_args = eval(f"dict{value[ref_key_end:]}", globals={}, locals={})
            else:
                ref_args = {}
            
            if ref_key in ctx:
            
                new_value = ctx[ref_key]
                new_value = deepcopy(new_value) if deepcopy_ else new_value
                
                if ref_args:
                    if isinstance(new_value, AliasDict) and "params" in new_value:
                        params:AliasDict = new_value["params"]
                        for k, v in ref_args.items():
                            params[k] = v
                    else:
                        raise ValueError(f"Reference does not has params: {ref_key}")
            
                dct[key] = new_value
            else:
                warnings.warn(f"Undefined reference: {ref_key}")
            

class AliasDict:

    def __init__(self, dct: dict, mods:dict, ctx: dict, ns:str, deepcopy_:bool):
        self.dct = dct
        _resolve_dollar(dct, mods, ctx, deepcopy_)
        self._build_recursive(mods, ctx, ns, deepcopy_)
        self._build_index(ctx, ns)
    
    def _build_index(self, ctx, ns):
        self.aliases = {}
        for key, value in self.dct.items():
            left_bracket = key.find("(")
            if left_bracket >= 0:
                right_bracket = key.rfind(")")
                if right_bracket < 0:
                    right_bracket = len(key)
                if right_bracket - left_bracket <= 1:
                    raise ValueError(f"Invalid alias: {key}")
                alias_key = key[left_bracket+1:right_bracket]
                if alias_key.startswith("$"):
                    base_class = alias_key[1:]
                    if not base_class in ctx:
                        raise ValueError(f"Undefined base class: {base_class}")
                    legacy = deepcopy(ctx[base_class])
                    if not isinstance(legacy, AliasDict):
                        raise ValueError(f"Base class is not a dict: {base_class}")
                    legacy.dct.setdefault("params", {})
                    if value == "pass":
                        pass
                    elif isinstance(value, AliasDict) and "params" in value:
                        for k, v in value['params'].items():
                            legacy.dct["params"][k] = v
                    else:
                        raise ValueError(f"Invalid alias value: {value}")
                    self.dct[key] = legacy
                else:
                    self.aliases[alias_key] = key
                    ctx[f"{ns}.{alias_key}"] = self.dct[key]
                official_key = key[:left_bracket]
                self.aliases[official_key] = key
                ctx[f"{ns}.{official_key}"] = self.dct[key]
            self.aliases[key] = key
            ctx[f"{ns}.{key}"] = self.dct[key]
    
    def _build_recursive(self, mods, ctx, ns:str, deepcopy_):
        for key, value in self.dct.items():
            if isinstance(value, dict):
                self.dct[key] = AliasDict(value, mods, ctx, f"{ns}.{key}", deepcopy_)
            else:
                self.dct[key] = value
    
    @property
    def official(self):
        official_dct = {}
        for key, value in self.dct.items():
            
            if isinstance(value, AliasDict):
                value = value.official
        
            left_bracket = key.find("(")
            if left_bracket >= 0:
                official_dct[key[:left_bracket]] = value
            else:
                official_dct[key] = value
        return official_dct
    
    def __getitem__(self, key):
        return self.dct[self.aliases[key]]
        
    def __setitem__(self, key, value):
        self.dct[self.aliases[key]] = value
    
    def __contains__(self, key):
        return key in self.aliases
    
    def __len__(self):
        return len(self.dct)
    
    def __iter__(self):
        return iter(self.dct)
    
    def __repr__(self):
        return f"AliasDict({self.dct})"
        
    def __rich_repr__(self):
        for key, value in self.dct.items():
            yield key, value
        
    def keys(self):
        return self.dct.keys()
    
    def values(self):
        return self.dct.values()
        
    def items(self):
        return self.dct.items()