import os

import gradio as gr
import wandb

from machine_learning.recommending.app import build_app_blocks, MovieMarkdownGenerator
from machine_learning.recommending.models.mf import MFRecommender
from machine_learning.recommending.movielens.data import MovieLens25m
from machine_learning.recommending.utils import wandb_timeit


project = "Recommending"
tmdb_api_token = os.environ["TMDB_API_TOKEN"]
lightning_class = MFRecommender
config = dict(
    artifact_name="my_mf_slim_movielens_25m:v0",
    movie_lens_25m_directory="ml-25m"
)

with wandb_timeit("wandb_init"):
    wandb.init(job_type="app", project=project, config=config)
artifact = wandb.use_artifact(config["artifact_name"])
with wandb_timeit("artifact_file"):
    checkpoint_path = artifact.file()
lightning_module = lightning_class.load_from_checkpoint(
    checkpoint_path, map_location="cpu"
)
model = lightning_module.model.eval()

movielens = MovieLens25m(path_to_movielens_folder=config["movie_lens_25m_directory"])
movie_markdown_generator = MovieMarkdownGenerator(
    movielens=movielens, tmdb_api_token=tmdb_api_token
)

with gr.Blocks() as app:
    build_app_blocks(
        recommender=model,
        movie_markdown_generator=movie_markdown_generator,
    )

app.launch()