import glob
import os
import pickle
from types import SimpleNamespace
from typing import List, Set

import streamlit as st
from dataclasses import dataclass
from natsort import natsorted


def clean_blocks(blocks):
    """
    Take streamlit placeholders and empty them.
    """
    for block in blocks:
        block.empty()


def setup_blocks(num_blocks):
    """
    Return num_blocks in the sidebar.
    """
    return [st.sidebar.empty() for _ in range(num_blocks)]


def get_user_history(username: str):
    """
    Load a user's history of interaction from session pickle files.
    """
    last_session_id = 0
    session_history = []
    if os.path.exists(f'data/{username}'):
        session_files = natsorted(list(glob.glob(f'data/{username}/session_*.pkl')))
        if session_files:
            last_session_id = int(session_files[-1].split("_")[1].split(".")[0])
            session_history = [pickle.load(open(file, 'rb')) for file in session_files]

    return SimpleNamespace(
        session_id=last_session_id + 1,
        session_history=session_history
    )


@dataclass
class Generation:
    # Model name
    model: str
    # Model checkpoint info
    checkpoint: dict
    # Model config info
    config: dict

    # Input text
    input: str
    # Generated output text
    output: str

    # Optional labels and other annotations
    labels: Set[str] = None
    annotations: List[str] = None


@dataclass
class Session:
    # Application mode
    mode: str
    # Session name and ID
    name: str
    id: int
    # Description of the session
    description: str

    # List of generations
    generations: List[Generation]

    # Favorited indices
    favorites: Set[int]

    # Attributes to label generated text
    attributes: list
    
    def to_pickle(self, path: str):
        with open(path, 'wb') as outfile:
            pickle.dump(self, outfile)
