import os

os.environ["PYTORCH_ENABLE_MPS_FALLBACK"] = "1"

import argparse
import requests
import shutil
import subprocess
import tempfile
import wandb
import wandb.apis.public

from typing import List, Optional

from huggingface_hub.hf_api import HfApi, upload_folder
from huggingface_hub.repocard import metadata_save
from pyvirtualdisplay.display import Display

from publish.markdown_format import EvalTableData, model_card_text
from runner.evaluate import EvalArgs, evaluate_model
from runner.env import make_eval_env
from shared.callbacks.eval_callback import evaluate
from wrappers.vec_episode_recorder import VecEpisodeRecorder


def publish(
    wandb_run_paths: List[str],
    wandb_report_url: str,
    huggingface_user: Optional[str] = None,
    huggingface_token: Optional[str] = None,
) -> None:
    virtual_display = Display(visible=False, size=(1400, 900))
    virtual_display.start()

    api = wandb.Api()
    runs = [api.run(rp) for rp in wandb_run_paths]
    algo = runs[0].config["algo"]
    env = runs[0].config["env"]
    evaluations = [
        evaluate_model(
            EvalArgs(
                algo,
                env,
                seed=r.config.get("seed", None),
                render=False,
                best=True,
                n_envs=None,
                n_episodes=10,
                no_print_returns=True,
                wandb_run_path="/".join(r.path),
            ),
            os.path.dirname(__file__),
        )
        for r in runs
    ]
    run_metadata = requests.get(runs[0].file("wandb-metadata.json").url).json()
    table_data = list(EvalTableData(r, e) for r, e in zip(runs, evaluations))
    best_eval = sorted(
        table_data, key=lambda d: d.evaluation.stats.score, reverse=True
    )[0]

    with tempfile.TemporaryDirectory() as tmpdirname:
        _, (policy, stats, config) = best_eval

        repo_name = config.model_name(include_seed=False)
        repo_dir_path = os.path.join(tmpdirname, repo_name)
        # Locally clone this repo to a temp directory
        subprocess.run(["git", "clone", ".", repo_dir_path])
        shutil.rmtree(os.path.join(repo_dir_path, ".git"))
        model_path = config.model_dir_path(best=True, downloaded=True)
        shutil.copytree(
            model_path,
            os.path.join(
                repo_dir_path, "saved_models", config.model_dir_name(best=True)
            ),
        )

        github_url = "https://github.com/sgoodfriend/rl-algo-impls"
        commit_hash = run_metadata.get("git", {}).get("commit", None)
        card_text = model_card_text(
            algo,
            env,
            github_url,
            commit_hash,
            wandb_report_url,
            table_data,
            best_eval,
        )
        readme_filepath = os.path.join(repo_dir_path, "README.md")
        os.remove(readme_filepath)
        with open(readme_filepath, "w") as f:
            f.write(card_text)

        metadata = {
            "library_name": "rl-algo-impls",
            "tags": [
                env,
                algo,
                "deep-reinforcement-learning",
                "reinforcement-learning",
            ],
            "model-index": [
                {
                    "name": algo,
                    "results": [
                        {
                            "metrics": [
                                {
                                    "type": "mean_reward",
                                    "value": str(stats.score),
                                    "name": "mean_reward",
                                }
                            ],
                            "task": {
                                "type": "reinforcement-learning",
                                "name": "reinforcement-learning",
                            },
                            "dataset": {
                                "name": env,
                                "type": env,
                            },
                        }
                    ],
                }
            ],
        }
        metadata_save(readme_filepath, metadata)

        video_env = VecEpisodeRecorder(
            make_eval_env(
                config,
                override_n_envs=1,
                normalize_load_path=model_path,
                **config.env_hyperparams,
            ),
            os.path.join(repo_dir_path, "replay"),
            max_video_length=3600,
        )
        evaluate(
            video_env,
            policy,
            1,
            deterministic=config.eval_params.get("deterministic", True),
        )

        api = HfApi()
        huggingface_user = huggingface_user or api.whoami()["name"]
        huggingface_repo = f"{huggingface_user}/{repo_name}"
        api.create_repo(
            token=huggingface_token,
            repo_id=huggingface_repo,
            private=False,
            exist_ok=True,
        )
        repo_url = upload_folder(
            repo_id=huggingface_repo,
            folder_path=repo_dir_path,
            path_in_repo="",
            commit_message=f"{algo.upper()} playing {env} from {github_url}/tree/{commit_hash}",
            token=huggingface_token,
        )
        print(f"Pushed model to the hub: {repo_url}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--wandb-run-paths",
        type=str,
        nargs="+",
        help="Run paths of the form entity/project/run_id",
    )
    parser.add_argument("--wandb-report-url", type=str, help="Link to WandB report")
    parser.add_argument(
        "--huggingface-user",
        type=str,
        help="Huggingface user or team to upload model cards",
        default=None,
    )
    args = parser.parse_args()
    print(args)
    publish(**vars(args))