from huggingface_hub import HfApi, ModelCard, CardData, upload_file
import json
import os

from huggingface_hub import login
login()

repo_id = "chrisjcc/utdg-maskableppo-policy"
model_path = "models/model_policy.zip"

api = HfApi()

# Create the repository first
try:
    api.create_repo(
        repo_id=repo_id,
        repo_type="model",
        exist_ok=True
    )
    print(f"✓ Repository {repo_id} created/verified")
except Exception as e:
    print(f"Repository creation: {e}")

# Upload model weights
print("Uploading model weights...")
upload_file(
    path_or_fileobj=model_path,
    path_in_repo=model_path,
    repo_id=repo_id,
    repo_type="model"
)
print("✓ Model uploaded")

# Create model card with CardData
card_data = CardData(
    language="en",
    license="mit",
    library_name="stable-baselines3",
    tags=["reinforcement-learning", "stable-baselines3", "gymnasium", "maskable-ppo"],
    datasets=["custom-utdg-env"],
    metrics=["episode_reward"],
)

card = ModelCard.from_template(
    card_data=card_data,
    model_id=repo_id,
)

card.text = """
# UTDG Maskable PPO Policy

This model is trained on the UTDG (Untitled Tower Defense Game) environment using Stable-Baselines3 MaskablePPO.

## Model Details

- **Algorithm**: MaskablePPO (Proximal Policy Optimization with invalid action masking)
- **Framework**: Stable-Baselines3
- **Environment**: Custom UTDG Gymnasium environment
- **Task**: Tower defense game AI agent

## Usage
```python
from huggingface_hub import hf_hub_download
from sb3_contrib import MaskablePPO

# Download the model
model_path = hf_hub_download(
    repo_id="chrisjcc/utdg-maskableppo-policy",
    filename="model_policy.zip"
)

# Load the model
model = MaskablePPO.load(model_path)

# Use for inference
# obs, info = env.reset()
# action, _states = model.predict(obs, action_masks=info["action_mask"])
```

## Training

The model was trained using reinforcement learning on the UTDG environment.
"""

card.save("README.md")
print("Uploading README...")
upload_file(
    path_or_fileobj="README.md",
    path_in_repo="README.md",
    repo_id=repo_id,
    repo_type="model"
)
print("✓ README uploaded")
print(f"\n✓ All done! View your model at: https://huggingface.co/{repo_id}")
