import txp.common.utils.bigquery_utils as bq_utils
import google.cloud.bigquery as bigquery
import txp.common.reports.section as sections
from google.oauth2 import service_account
import plotly.graph_objects as go
from PIL import Image
from typing import Dict
import pandas as pd
import json
import logging
from txp.common.config import settings
log = logging.getLogger(__name__)
log.setLevel(settings.txp.general_log_level)


def get_last_section_plot(
    tenant_id: str,
    section_type: sections.ReportSection,
    bq_client: bigquery.Client
) -> Image:
    """Provides access to the most recent static image plot generated by a
    section.

    Args:
         tenant_id: The tenant id is required to perform specific queries.
         section_type: The class reference to the section to be requested.
         bq_client: bigquery.Client instance.

    Returns:
        The PIL.Image if found, None otherwise.
        None in case of internal exception.
    """
    # perform bigquery query
    log.info(f"Requesting {str(type(section_type))} static image for tenant {tenant_id}")
    section_df: pd.DataFrame = bq_utils.get_last_sections_from_bq(
        'reports.sections',
        tenant_id,
        section_type.__name__,
        bq_client
    )

    if not section_df.size:
        log.info(f"No static image found for {str(type(section_type))}")
        return None

    # Gets the section registry by loading the DB dict
    section_registry: Dict = section_df.to_dict(orient='index')[0]

    # Get the image from the registry
    image_bytes = json.loads(section_registry['data'])['image']
    img_pil = section_type.deserialize_image(image_bytes)

    return img_pil


def get_last_section_dinamic_plot(
    tenant_id: str,
    section_type: sections.ReportSection,
    bq_client: bigquery.Client
) -> go.Figure:
    """Provides access to the most recent plotly plot generated by a
    section.

    Args:
         tenant_id: The tenant id is required to perform specific queries.
         section_type: The class reference to the section to be requested.
         bq_client: bigquery.Client instance.

    Returns:
        The plotly Figure if found, None otherwise.
        None in case of internal exception.
    """
    # perform bigquery query
    # perform bigquery query
    log.info(f"Requesting {str(type(section_type))} dynamic plot for tenant {tenant_id}")

    section_df: pd.DataFrame = bq_utils.get_last_sections_from_bq(
        'reports.sections',
        tenant_id,
        section_type.__name__,
        bq_client
    )

    if not section_df.size:
        log.info(f"No static image found for {str(type(section_type))}")
        return None

    # Gets the section dinamic plotly to display with streamlit
    sc = section_type.load(section_df.to_dict(orient='index')[0])
    fig = sc.get_dynamic_plot()

    return fig


if __name__ == '__main__':
    """Main for fast testing"""
    from txp.common.reports.mock_double_line import MockDoubleLine
    from txp.common.reports.forecasting import StatesForecasting
    from txp.common.reports.anomaly_detection import AnomalyDetectionStates

    CREDENTIALS_PATH = "/Users/marcosandoval/Work/Tranxpert/txp/src/txp/common/credentials/pub_sub_to_bigquery_credentials.json"
    credentials = service_account.Credentials.from_service_account_file(CREDENTIALS_PATH)
    client = bigquery.Client(credentials=credentials, project=credentials.project_id)

    plot = get_last_section_dinamic_plot(
        'labshowroom-001',
        AnomalyDetectionStates,
        client
    )

    plot.show()
