gisaf-backend/src/gisaf/models/dashboard.py

from io import BytesIO
from pickle import loads
from pathlib import Path
from datetime import datetime
import logging
# from typing import Any

from matplotlib.figure import Figure
from sqlmodel import Field, Relationship, String, JSON
from pydantic import BaseModel
import pandas as pd

from gisaf.config import conf
from gisaf.models.metadata import gisaf
from gisaf.models.models_base import Model
from gisaf.models.misc import NotADataframeError

import matplotlib.pyplot as plt

logger = logging.getLogger(__name__)



class DashboardPageSource(Model, table=True):
    __tablename__ = 'dashboard_page_source' # type: ignore
    __table_args__ = gisaf.table_args

    id: str = Field(primary_key=True)
    name: str


class DashboardPageCommon:
    """
    Base class for DashboardPage and DashboardPageSection, where some methods
    are common, eg. attachments
    """
    name: str
    df: bytes
    plot: bytes
    #plot: dict[str, Any] | None = Field(sa_type=JSON(none_as_null=True)) # type: ignore
    attachment: str | None
    html: str | None = None

    def ensure_dir_exists(self):
        raise NotImplementedError()

    def get_plot_file_path(self) -> Path:
        raise NotImplementedError()

    def get_attachment_file_name(self) -> str:
        raise NotImplementedError()

    def get_plot_file_name(self):
        raise NotImplementedError()

    def save_plot(self, plot: plt.Axes | plt.Figure): # type: ignore
        """
        Render the matplotlib plot (or figure) and save it in the filesystem
        :param plot: matplotlib plot or figure...
        :return:
        """
        self.ensure_dir_exists()

        ## Different types of figures supported
        fig: Figure | None = None
        if isinstance(plot, plt.Axes): # type: ignore
            fig = plot.figure # type: ignore
        elif isinstance(plot, plt.Figure): # type: ignore
            fig = plot
        if fig:
            fig.savefig(self.get_plot_file_path(), bbox_inches='tight')
            plt.close(fig)

        if plot and not fig:
            logger.warning('Cannot save dashboard attachment (unknown attachment type)')
            return

        #logger.info(f'Saved attachment of dashboard page {self.group}/{self.name} '
                    #f'in {self.get_attachment_file_name()}')
        return self.get_plot_file_name()

    def save_attachment(self, attached, name=None) -> str | None:
        """
        Save the attachment in the filesystem
        :param attached: matplotlib plot or figure...
        :return:
        """
        if not self.attachment:
            ## Not set yet (creation)
            self.attachment = f'{self.name}.png'

        self.ensure_dir_exists()

        ## Different types of figures supported
        fig: Figure | None = None
        if plt:
            if isinstance(attached, plt.Axes): # type: ignore
                fig = attached.figure # type: ignore
            elif isinstance(attached, plt.Figure): # type: ignore
                fig = attached
            if fig:
                fig.savefig(self.get_attachment_file_name(), bbox_inches='tight')
                plt.close(fig)

        if attached and not fig:
            logger.warning('Cannot save dashboard attachment (unknown attachment type)')
            return None

        #logger.info(f'Saved attachment of dashboard page {self.group}/{self.name} '
                    #f'in {self.get_attachment_file_name()}')
        return self.attachment

    def get_page_df(self):
        """
        Get the dataframe of the page
        """
        if not self.df:
            return
        try:
            return pd.read_pickle(BytesIO(self.df), compression=None)
        except KeyError:
            raise NotADataframeError()

    def get_plot(self):
        if self.plot is not None:
            return loads(self.plot)


class DashboardPageMetaData(BaseModel):
    name: str
    group: str
    description: str
    viewable_role: str | None = None


class DashboardPage(Model, DashboardPageCommon, DashboardPageMetaData, table=True):
    __tablename__ = 'dashboard_page' # type: ignore
    __table_args__ = gisaf.table_args

    class Admin:
        menu = 'Dashboard'

    id: int = Field(primary_key=True)
    time: datetime | None = Field(default_factory=datetime.now)
    notebook: str | None = None
    source_id: int | None = Field(foreign_key=gisaf.table('dashboard_page_source.id'))
    expanded_panes: str | None
    source: DashboardPageSource = Relationship()
    sections: list['DashboardPageSection'] = Relationship()

    def __str__(self):
        return f'{self.group:s}/{self.name:s}'

    def __repr__(self):
        return f'<models.DashboardPage {self.group:s}/{self.name:s}>'

    @classmethod
    def selectinload(cls):
        return [cls.source]

    def ensure_dir_exists(self):
        """
        Make sure the directory exists, before saving the file
        """
        dir_name = Path(conf.dashboard.base_storage_dir) / self.group
        dir_name.mkdir(exist_ok=True)
        return dir_name

    def get_attachment_file_name(self):
        """
        Get the file name of the attachment
        :return:
        """
        if self.attachment:
            base_dir = Path(conf.dashboard.base_storage_dir)
            if base_dir:
                return base_dir/self.group/self.attachment
            else:
                raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')

    def get_attachment_url(self):
        ## Serve through web front-end (nginx static file)
        if not self.attachment:
            return
        base_storage_url = conf.dashboard.base_storage_url
        if not base_storage_url:
            base_storage_url = '/dashboard-attachment/'
        return f'{base_storage_url}{self.group}/{self.attachment}'

    def get_notebook_url(self):
        if self.notebook:
            base_url = conf.dashboard.base_source_url
            if base_url:
                return f'{base_url}{self.notebook}'
            else:
                logger.debug('Notebook: no base_url in gisaf config')


class DashboardPageSection(Model, DashboardPageCommon, table=True):
    __tablename__ = 'dashboard_page_section' # type: ignore
    __table_args__ = gisaf.table_args

    class Admin:
        menu = 'Dashboard'

    id: str = Field(primary_key=True)
    name: str
    dashboard_page_id: int = Field(foreign_key=gisaf.table('dashboard_page.id'))
    dashboard_page: DashboardPage = Relationship(back_populates='sections')

    description: str

    def __str__(self):
        return f'{self.name} for dashboard page #{self.dashboard_page_id}'

    def __repr__(self):
        return f'<models.DashboardPageSection #{self.id}>'

    @classmethod
    def selectinload(cls):
        return [cls.dashboard_page]

    def get_plot_url(self):
        ## Serve through web front-end (nginx static file)
        if not self.plot:
            return
        return conf.dashboard.base_storage_url \
            + self.dashboard_page.group + '/' \
            + self.dashboard_page.name + '/' \
            + self.name + '.png'

    def ensure_dir_exists(self):
        """
        Make sure the directory exists, before saving the file
        """
        dir_name = Path(conf.dashboard.base_storage_dir) / self.page.group / self.page.name
        dir_name.mkdir(exist_ok=True)
        return dir_name

    def get_attachment_file_name(self):
        """
        Get the file name of the attachment
        :return:
        """
        if not self.attachment:
            return
        base_dir = Path(conf.dashboard.base_storage_dir)
        if not base_dir:
            raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')
        return base_dir/self.page.group/self.page.name/self.attachment

    def get_plot_file_name(self):
        return f'{self.name}.png'

    def get_plot_file_path(self):
        """
        Get the file name of the plot
        :return:
        """
        if not self.plot:
            return
        base_dir = Path(conf.dashboard.base_storage_dir)
        if not base_dir:
            raise UserWarning('Cannot save attachment: no notebook/base_storage_dir in gisaf config')
        return base_dir/self.page.group/self.page.name/self.get_plot_file_name()


class Widget(Model, table=True):
    __tablename__ = 'widget' # type: ignore
    __table_args__ = gisaf.table_args
    ## CREATE TABLE gisaf.widget (name char(50) not null PRIMARY KEY, title varchar, subtitle varchar, notebook varchar, content varchar, time timestamp);
    name: str = Field(primary_key=True, sa_type=String(50)) # type: ignore
    title: str
    subtitle: str
    content: str
    time: datetime
    notebook: str

    class Admin:
        menu = 'Dashboard'


class DashboardSection(BaseModel):
    name: str
    plot: str


class Dashboard(BaseModel):
    name: str
    group: str
    description: str
    time: datetime | None = Field(default_factory=datetime.now)
    html: str | None = None
    attachment: str | None = None
    dfData: list = []
    plotData: str | None = None
    notebook: str | None = None
    expandedPanes: list[str] | None = None
    sections: list[DashboardSection] | None = None


class DashboardGroup(BaseModel):
    name: str
    pages: list[DashboardPageMetaData]


class DashboardHome(BaseModel):
    title: str
    content: str
    footer: str