import os
import json
import torch
from easydict import EasyDict as edict
from anigen import models

def load_model_from_path(path, model_name_in_config=None, device='cuda', use_ema=False):
    if os.path.isdir(path):
        config_path = os.path.join(path, 'config.json')
        if not os.path.exists(config_path):
            raise ValueError(f"Config file not found in {path}")
        with open(config_path, 'r') as f:
            config = json.load(f)
        config = edict(config)
        
        ckpt_dir = os.path.join(path, 'ckpts')
        if not os.path.exists(ckpt_dir):
             raise ValueError(f"Checkpoints directory not found in {path}")
        
        files = os.listdir(ckpt_dir)
        pt_files = [f for f in files if f.endswith('.pt')]
        if not pt_files:
            raise ValueError(f"No .pt files found in {ckpt_dir}")
        
        def get_step(name):
            try:
                return int(name.split('step')[-1].split('.')[0])
            except:
                return -1
        
        # Filter for EMA if requested
        if use_ema:
            ema_files = [f for f in pt_files if 'ema' in f]
            if ema_files:
                pt_files = ema_files
                print("Selected EMA checkpoint.")
            else:
                print("Warning: EMA checkpoint requested but not found. Falling back to regular checkpoint.")
                pt_files = [f for f in pt_files if 'ema' not in f and 'misc' not in f]
        else:
            # Exclude 'misc' checkpoints which contain optimizer state, not model weights
            non_ema_files = [f for f in pt_files if 'ema' not in f and 'misc' not in f]
            if non_ema_files:
                pt_files = non_ema_files
                print("Selected regular checkpoint.")
            else:
                print("Warning: Regular checkpoint not found. Falling back to EMA checkpoint.")
                pt_files = [f for f in pt_files if 'ema' in f]
        
        pt_files.sort(key=get_step, reverse=True)
        ckpt_path = os.path.join(ckpt_dir, pt_files[0])
        print(f"Loading checkpoint: {ckpt_path}")
        
        if model_name_in_config:
            model_config = config.models[model_name_in_config]
        else:
            keys = list(config.models.keys())
            # Heuristic: prefer 'denoiser' or 'flow_model'
            if 'denoiser' in keys:
                model_config = config.models['denoiser']
            elif len(keys) == 1:
                model_config = config.models[keys[0]]
            else:
                raise ValueError(f"Multiple models in config {keys}, please specify model_name_in_config")

        model = getattr(models, model_config.name)(**model_config.args)
        state_dict = torch.load(ckpt_path, map_location='cpu')
        
        if list(state_dict.keys())[0].startswith('module.'):
            state_dict = {k[7:]: v for k, v in state_dict.items()}
            
        model.load_state_dict(state_dict, strict=False)
        model.to(device)
        model.eval()
        return model, config
    else:
        raise ValueError("Please provide a directory containing config.json and ckpts/")

def load_decoder(path, ckpt_name, device):
    if not os.path.exists(path):
        raise ValueError(f"Decoder path not found: {path}")
    
    config_path = os.path.join(path, 'config.json')
    if not os.path.exists(config_path):
        raise ValueError(f"Config file not found in {path}")
        
    with open(config_path, 'r') as f:
        cfg = json.load(f)
    
    if 'models' not in cfg or 'decoder' not in cfg['models']:
        raise ValueError(f"Config at {path} does not have ['models']['decoder']")
        
    model_cfg = cfg['models']['decoder']
    decoder = getattr(models, model_cfg['name'])(**model_cfg['args'])
    
    ckpt_path = os.path.join(path, 'ckpts', f'decoder_{ckpt_name}.pt')
    if not os.path.exists(ckpt_path):
        # Fallback to just ckpt_name if decoder_ prefix not found
        ckpt_path = os.path.join(path, 'ckpts', f'{ckpt_name}.pt')
        if not os.path.exists(ckpt_path):
             raise ValueError(f"Checkpoint not found: {ckpt_path}")
            
    print(f"Loading decoder from {ckpt_path}")
    state_dict = torch.load(ckpt_path, map_location='cpu')
    if list(state_dict.keys())[0].startswith('module.'):
        state_dict = {k[7:]: v for k, v in state_dict.items()}
    decoder.load_state_dict(state_dict, strict=False)
    decoder.to(device).eval()
    return decoder