# HF space creator starting from an sklearn model

import base64
import glob
import io
import json
import os
import pickle
import re
import shutil
from pathlib import Path
from tempfile import mkdtemp

import pandas as pd
import sklearn
import streamlit as st
from sklearn.base import BaseEstimator

import skops.io as sio
from skops import card, hub_utils

st.set_page_config(layout="wide")
st.title("Skops space creator for sklearn")


PLACEHOLDER = "[More Information Needed]"
PLOT_PREFIX = "__plot__:"
custom_sections: dict[str, str] = {}
tmp_repo = Path(mkdtemp(prefix="skops-"))
left_col, right_col = st.columns([1, 2])

# a hacky way to "persist" custom sections
CUSTOM_SECTIONS_CACHE_FILE = ".custom-sections.json"


def _clear_custom_section_cache():
    with open(CUSTOM_SECTIONS_CACHE_FILE, "w") as f:
        f.write("")


def _load_custom_section_cache():
    global custom_sections

    # in case file doesn't exist yet, create it
    if not os.path.exists(CUSTOM_SECTIONS_CACHE_FILE):
        Path(CUSTOM_SECTIONS_CACHE_FILE).touch()

    with open(CUSTOM_SECTIONS_CACHE_FILE, "r") as f:
        try:
            custom_sections = json.load(f)
        except ValueError:
            pass


def _write_custom_section_cache():
    with open(CUSTOM_SECTIONS_CACHE_FILE, "w") as f:
        json.dump(custom_sections, f)


def _remove_custom_section(key):
    del custom_sections[key]
    _write_custom_section_cache()


def _clear_repo(path):
    for file_path in glob.glob(str(Path(path) / "*")):
        if os.path.isfile(file_path) or os.path.islink(file_path):
            os.unlink(file_path)
        elif os.path.isdir(file_path):
            shutil.rmtree(file_path)


def _write_plot(plot_name, plot_file):
    with open(plot_name, "wb") as f:
        f.write(plot_file)


def init_repo():
    _clear_repo(tmp_repo)

    try:
        file_name = Path(mkdtemp(prefix="skops-")) / "model.skops"
        sio.dump(model, file_name)
        hub_utils.init(
            model=file_name,
            dst=tmp_repo,
            task=task,
            data=data,
            requirements=requirements,
        )
    except Exception as exc:
        print("Uh oh, something went wrong when initializing the repo:", exc)


def load_model():
    if model_file is None:
        return

    bytes_data = model_file.getvalue()
    model = pickle.loads(bytes_data)
    assert isinstance(model, BaseEstimator), "model must be an sklearn model"
    return model


def load_data():
    if data_file is None:
        return

    bytes_data = io.BytesIO(data_file.getvalue())
    df = pd.read_csv(bytes_data)
    return df


def _parse_metrics(metrics):
    metrics_table = {}
    for line in metrics.splitlines():
        line = line.strip()
        name, _, val = line.partition("=")
        try:
            # try to coerce to float but don't error if it fails
            val = float(val.strip())
        except ValueError:
            pass
        metrics_table[name.strip()] = val
    return metrics_table


def _create_model_card():
    if model is None or data is None:
        st.text("*some data is missing to render the model card*")
        return

    init_repo()
    metadata = card.metadata_from_config(tmp_repo)
    model_card = card.Card(model=model, metadata=metadata)

    if model_description:
        model_card.add(**{"Model description": model_description})

    if intended_uses:
        model_card.add(
            **{"Model description/Intended uses & limitations": intended_uses}
        )

    if metrics:
        metrics_table = _parse_metrics(metrics)
        model_card.add_metrics(**metrics_table)

    if authors:
        model_card.add(**{"Model Card Authors": authors})

    if contact:
        model_card.add(**{"Model Card Contact": contact})

    if citation:
        model_card.add(**{"Citation": citation})

    if custom_sections:
        for key, val in custom_sections.items():
            if not key:
                continue

            if key.startswith(PLOT_PREFIX):
                key = key[len(PLOT_PREFIX):]
                model_card.add_plot(**{key: val})
            else:
                model_card.add(**{key: val})

    return model_card


def _process_card_for_rendering(rendered: str) -> tuple[str, str]:
    idx = rendered[1:].index("\n---") + 1
    metadata = rendered[3:idx]
    rendered = rendered[idx + 4 :]  # noqa: E203

    # below is a hack to display the images in streamlit
    # https://discuss.streamlit.io/t/image-in-markdown/13274/10 The problem is

    # that streamlit does not display images in markdown, so we need to replace
    # them with html. However, we only want that in the rendered markdown, not
    # in the card that is produced for the hub
    def markdown_images(markdown):
        # example image markdown:
        # ![Test image](images/test.png "Alternate text")
        images = re.findall(
            r'(!\[(?P<image_title>[^\]]+)\]\((?P<image_path>[^\)"\s]+)\s*([^\)]*)\))',
            markdown
        )
        return images

    def img_to_bytes(img_path):
        img_bytes = Path(img_path).read_bytes()
        encoded = base64.b64encode(img_bytes).decode()
        return encoded

    def img_to_html(img_path, img_alt):
        img_format = img_path.split(".")[-1]
        img_html = (
            f'<img src="data:image/{img_format.lower()};'
            f'base64,{img_to_bytes(img_path)}" '
            f'alt="{img_alt}" '
            'style="max-width: 100%;">'
        )
        return img_html

    def markdown_insert_images(markdown):
        images = markdown_images(markdown)

        for image in images:
            image_markdown = image[0]
            image_alt = image[1]
            image_path = image[2]
            markdown = markdown.replace(image_markdown, img_to_html(image_path, image_alt))
        return markdown

    rendered_with_img = markdown_insert_images(rendered)
    return metadata, rendered_with_img


def display_model_card():
    model_card = _create_model_card()
    if not model_card:
        return

    rendered = model_card.render()
    metadata, rendered = _process_card_for_rendering(rendered)
    # idx = rendered[1:].index("\n---") + 1
    # metadata = rendered[3:idx]
    # rendered = rendered[idx + 4 :]  # noqa: E203

    with right_col:
        # strip metadata
        with st.expander("show metadata"):
            st.text(metadata)
        st.markdown(rendered, unsafe_allow_html=True)


def download_model_card():
    model_card = _create_model_card()
    if model_card is not None:
        return model_card.render()
    return ""


def add_custom_section():
    # this is required to "refresh" these variables...
    global section_name, section_content
    section_name = st.session_state.key_section_name
    section_content = st.session_state.key_section_content

    if not section_name or not section_content:
        return

    custom_sections[section_name] = section_content
    _write_custom_section_cache()


def add_custom_plot():
    # this is required to "refresh" these variables...
    global section_name, section_content
    plot_name = st.session_state.key_plot_name
    plot_file = st.session_state.key_plot_file

    if not plot_name or not plot_file:
        return

    # store plot in temp repo
    file_name = plot_file.name.replace(" ", "_")
    file_path = str(tmp_repo / file_name)
    with open(file_path, "wb") as f:
        f.write(plot_file.getvalue())

    custom_sections[str(PLOT_PREFIX + plot_name)] = file_path
    _write_custom_section_cache()


with left_col:
    # This contains every element required to edit the model card
    model = None
    data = None
    section_name = None
    section_content = None

    model_file = st.file_uploader("Upload a model*", on_change=load_model)
    data_file = st.file_uploader(
        "Upload X data (csv)*", type=["csv"], on_change=load_data
    )

    task = st.selectbox(
        label="Choose the task type*",
        options=[
            "tabular-classification",
            "tabular-regression",
            "text-classification",
            "text-regression",
        ],
        on_change=init_repo,
    )

    requirements = st.text_input(
        label="Requirements*",
        value=[f"scikit-learn=={sklearn.__version__}\n"],
        on_change=init_repo,
    )

    if model_file is not None:
        model = load_model()

    if data_file is not None:
        data = load_data()

    if model is not None and data is not None:
        init_repo()

    model_description = st.text_input("Model description", value=PLACEHOLDER)
    intended_uses = st.text_area(
        "Intended uses & limitations", height=2, value=PLACEHOLDER
    )
    metrics = st.text_area("Metrics (e.g. 'accuracy = 0.95'), one metric per line")
    authors = st.text_area(
        "Authors",
        value="This model card is written by following authors:\n\n" + PLACEHOLDER,
    )
    contact = st.text_area(
        "Contact",
        value="You can contact the model card authors through following channels:\n\n"
        + PLACEHOLDER,
    )
    citation = st.text_area(
        "Citation",
        value="Below you can find information related to citation.\n\nBibTex:\n\n```\n"
        + PLACEHOLDER
        + "\n```",
        height=5,
    )

    with st.form("custom-section", clear_on_submit=True):
        section_name = st.text_input(
            "Section name (use '/' for subsections, e.g. 'Model description/My new"
            " section')",
            key="key_section_name",
        )
        section_content = st.text_area(
            "Content of the new section", key="key_section_content"
        )
        submit_new_section = st.form_submit_button(
            "Create new section", on_click=add_custom_section
        )

    with st.form("custom-plots", clear_on_submit=True):
        plot_name = st.text_input(
            "Section name (use '/' for subsections, e.g. 'Model description/My new"
            " plot')",
            key="key_plot_name",
        )
        plot_file = st.file_uploader(
            "Upload a figure*", key="key_plot_file"
        )

        submit_new_plot = st.form_submit_button(
            "Add plot", on_click=add_custom_plot
        )

    _load_custom_section_cache()
    for key in custom_sections:
        if not key:
            continue

        if key.startswith(PLOT_PREFIX):
            st.button(
                f"Remove plot '{key[len(PLOT_PREFIX):]}'", on_click=_remove_custom_section, args=(key,)
            )
        else:
            st.button(
                f"Remove section '{key}'", on_click=_remove_custom_section, args=(key,)
            )

    if custom_sections:
        st.button(
            f"Remove all ({len(custom_sections)}) custom elements",
            on_click=_clear_custom_section_cache,
        )


with right_col:
    # this contains the rendered model card
    st.button(label="Render model card", on_click=display_model_card)
    rendered = download_model_card()
    if rendered:
        st.download_button(label="Download model card (markdown format)", data=rendered)