from huggingface_hub import hf_hub_download
from safetensors.torch import load_model as st_load_model
import os, torch
from agent.sac import Actor, device

# Configuration
REPO_ID = "S-Mithilesh/Energizer-V1"
FILENAME = "model.safetensors"
DATASET_FILENAME = "sample_dataset.pkl"
STATE_DIM = 11

def load_trained_model(repo_id=REPO_ID, filename=FILENAME, use_local=None):
    """
    Loads model weights. 
    Priority: 1. Forced local (if use_local is set) 2. HF Hub 3. Local safetensors 4. Local .pt
    """
    model = Actor(STATE_DIM).to(device)
    token = os.environ.get("HF_TOKEN") # Automatic support for private spaces
    
    # 1. Explicit Local Override
    if use_local and os.path.exists(use_local):
        print(f"Loading forced local weights from {use_local}...")
        try:
            st_load_model(model, use_local)
            print("[SUCCESS] Successfully loaded forced local weights.")
            return model.eval()
        except Exception as e:
            print(f"[ERROR] Failed to load forced local weights: {e}")

    # 2. Try Hugging Face Hub (Preferred)
    try:
        print(f"Fetching weights from HF Hub: {repo_id}/{filename}...")
        # Download file from HF Hub (cached locally)
        path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)
        st_load_model(model, path)
        print("[SUCCESS] Model weights loaded successfully from HF Hub.")
        return model.eval()
    except Exception as e:
        print(f"[WARNING] Hub fetch failed: {e}")

    # 3. Fallback to local safetensors (if exists in repo)
    local_path = "models/model.safetensors"
    if os.path.exists(local_path):
        print(f"[FALLBACK] Loading local weights from {local_path}...")
        try:
            st_load_model(model, local_path)
            print("[SUCCESS] Loaded local safetensors fallback.")
            return model.eval()
        except Exception as e:
            print(f"[WARNING] Local safetensors load failed: {e}")

    # 4. Final fallback for original .pt file
    pt_path = "models/policy_sac.pt"
    if os.path.exists(pt_path):
        print(f"[FALLBACK] Loading local .pt checkpoint...")
        try:
            ckpt = torch.load(pt_path, map_location=device, weights_only=False)
            sd = ckpt['actor_state'] if isinstance(ckpt, dict) and 'actor_state' in ckpt else ckpt
            model.load_state_dict(sd)
            print("[SUCCESS] Loaded local .pt fallback.")
        except Exception as e:
            print(f"[ERROR] Local .pt load failed: {e}")
    else:
        print("[ERROR] No weights found. Model will have random performance.")

    model.eval()
    return model

def ensure_dataset_exists(repo_id=REPO_ID, filename=DATASET_FILENAME, target_dir="data"):
    """
    Ensures the sample dataset is available locally.
    Priority: 1. Local data/ folder 2. HF Hub download
    """
    target_path = os.path.join(target_dir, filename)
    if os.path.exists(target_path):
        return target_path
    
    token = os.environ.get("HF_TOKEN")
    try:
        print(f"Fetching sample dataset from HF Hub: {repo_id}/{filename}")
        os.makedirs(target_dir, exist_ok=True)
        # Download and cache
        path = hf_hub_download(repo_id=repo_id, filename=filename, token=token)
        # Move/Link to the expected data directory for environment consistency
        import shutil
        shutil.copy(path, target_path)
        print(f"[SUCCESS] Sample dataset localized to {target_path}")
        return target_path
    except Exception as e:
        print(f"[WARNING] Could not fetch sample dataset from Hub: {e}")
        return None

if __name__ == "__main__":
    # Test loading
    m = load_trained_model()
    d = ensure_dataset_exists()
    print("Test complete.")