from __future__ import annotations

import datetime as dt
import json
import shutil
from pathlib import Path

import gymnasium as gym
import imageio
import numpy as np
from huggingface_hub import HfApi
from huggingface_hub.errors import HfHubHTTPError
from huggingface_hub.repocard import metadata_eval_result, metadata_save
from stable_baselines3 import PPO
from stable_baselines3.common.env_util import make_vec_env
from stable_baselines3.common.evaluation import evaluate_policy
from stable_baselines3.common.monitor import Monitor


USERNAME = "Sami94"
STUDENT_NAME = "Sami Chellia"
ENV_ID = "LunarLander-v2"
MODEL_NAME = "ppo-LunarLander-v2"
REPO_ID = f"{USERNAME}/{MODEL_NAME}"
OUTPUT_DIR = Path("artifacts/unit1") / MODEL_NAME


def evaluate(model: PPO, n_eval_episodes: int = 10) -> tuple[float, float]:
    eval_env = Monitor(gym.make(ENV_ID, render_mode="rgb_array"))
    mean_reward, std_reward = evaluate_policy(
        model,
        eval_env,
        n_eval_episodes=n_eval_episodes,
        deterministic=True,
    )
    eval_env.close()
    return float(mean_reward), float(std_reward)


def record_video(model: PPO, out_path: Path, max_steps: int = 1000) -> None:
    env = gym.make(ENV_ID, render_mode="rgb_array")
    obs, _ = env.reset(seed=42)
    frames = [env.render()]
    for _ in range(max_steps):
        action, _ = model.predict(obs, deterministic=True)
        obs, _, terminated, truncated, _ = env.step(action)
        frames.append(env.render())
        if terminated or truncated:
            break
    env.close()
    imageio.mimsave(out_path, [np.asarray(frame) for frame in frames], fps=30)


def save_artifacts(model: PPO, mean_reward: float, std_reward: float, timesteps: int) -> None:
    if OUTPUT_DIR.exists():
        shutil.rmtree(OUTPUT_DIR)
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)

    model.save(OUTPUT_DIR / MODEL_NAME)
    record_video(model, OUTPUT_DIR / "replay.mp4")

    results = {
        "env_id": ENV_ID,
        "mean_reward": mean_reward,
        "std_reward": std_reward,
        "n_eval_episodes": 10,
        "total_timesteps": timesteps,
        "eval_datetime": dt.datetime.now().isoformat(),
        "student": STUDENT_NAME,
        "hf_username": USERNAME,
    }
    (OUTPUT_DIR / "results.json").write_text(json.dumps(results, indent=2), encoding="utf-8")

    metadata = {
        "tags": [
            ENV_ID,
            "ppo",
            "stable-baselines3",
            "reinforcement-learning",
            "huggingface-deep-rl-course",
        ]
    }
    eval_metadata = metadata_eval_result(
        model_pretty_name=MODEL_NAME,
        task_pretty_name="reinforcement-learning",
        task_id="reinforcement-learning",
        metrics_pretty_name="mean_reward",
        metrics_id="mean_reward",
        metrics_value=f"{mean_reward:.2f} +/- {std_reward:.2f}",
        dataset_pretty_name=ENV_ID,
        dataset_id=ENV_ID,
    )
    metadata = {**metadata, **eval_metadata}

    readme_path = OUTPUT_DIR / "README.md"
    readme_path.write_text(
        f"""# PPO Agent for {ENV_ID}

Student: {STUDENT_NAME}
Hugging Face username: {USERNAME}

This repository contains a PPO agent trained for the Hugging Face Deep RL course.

Mean reward: {mean_reward:.2f} +/- {std_reward:.2f}
Total timesteps: {timesteps}

```python
from huggingface_hub import hf_hub_download
from stable_baselines3 import PPO

model_path = hf_hub_download(repo_id="{REPO_ID}", filename="{MODEL_NAME}.zip")
model = PPO.load(model_path)
```
""",
        encoding="utf-8",
    )
    metadata_save(readme_path, metadata)


def push_artifacts() -> str:
    api = HfApi()
    try:
        repo_url = api.create_repo(repo_id=REPO_ID, repo_type="model", exist_ok=True)
        api.upload_folder(repo_id=REPO_ID, repo_type="model", folder_path=OUTPUT_DIR, path_in_repo=".")
        return str(repo_url)
    except HfHubHTTPError as exc:
        print(f"Hub push failed for {REPO_ID}: {exc}")
        print(f"Local artifacts saved in {OUTPUT_DIR.resolve()}")
        return f"LOCAL_ONLY:{OUTPUT_DIR.resolve()}"


def main() -> None:
    env = make_vec_env(ENV_ID, n_envs=16)
    model = PPO(
        policy="MlpPolicy",
        env=env,
        n_steps=1024,
        batch_size=64,
        n_epochs=4,
        gamma=0.999,
        gae_lambda=0.98,
        ent_coef=0.01,
        verbose=1,
    )

    total_timesteps = 0
    best: tuple[float, float] | None = None
    for chunk in [200_000, 200_000, 200_000, 200_000, 200_000]:
        model.learn(total_timesteps=chunk, reset_num_timesteps=False)
        total_timesteps += chunk
        mean_reward, std_reward = evaluate(model)
        best = (mean_reward, std_reward)
        print(f"Evaluation after {total_timesteps} timesteps: {mean_reward:.2f} +/- {std_reward:.2f}")
        save_artifacts(model, mean_reward, std_reward, total_timesteps)
        if mean_reward >= 200:
            print("Certification threshold reached for Unit 1.")
            break

    env.close()
    if best is None:
        raise RuntimeError("Training finished without evaluation.")

    url = push_artifacts()
    (OUTPUT_DIR / "pushed_repo.json").write_text(
        json.dumps({"repo_id": REPO_ID, "url": url}, indent=2),
        encoding="utf-8",
    )
    print(json.dumps({"repo_id": REPO_ID, "url": url, "mean_reward": best[0], "std_reward": best[1]}, indent=2))


if __name__ == "__main__":
    main()