# Copyright (c) 2023, Tony Liu
#
# Licensed under the Apache License, Version 2.0 (the "License");
# Use, reproduction and distribution of this software in source and 
# binary forms, with or without modification, are permitted provided that
# the License terms and conditions are met; you may not use this file
# except in compliance with the License. See the LICENSE file for details.

"""
This module contains StreamlitApp class representing an application
"""

__author__ = "Tony Liu"
__email__ = "tony.liu@yahoo.com"
__license__ = "Apache License 2.0"
__version__ = "0.2.0"


import os
from os.path import exists
import sys
from typing import Optional, Union, Dict, Callable
import logging
import toml
from functools import wraps
import streamlit as st

from snowflake_ai.common import AppConfig
from snowflake_ai.apps import AppPage
from snowflake_ai.common import OAuthConnect
from snowflake_ai.connect import AuthCodeConnect, ConnectManager



class StreamlitApp(AppConfig):
    """
    This class represents a Steamlit application's configurations and 
    its corresponding application. Since it is subclass of AppConfig,
    it has the following config directory bootstrapping precedence:

    1) input custom directory 
    2) snowflake_ai/conf subdir under current directory 
    3) snowflake_ai/conf subdir under user_home directory
    4) current directory
    5) home directory
    6) conf subdir under snowflake_ai library installation root dir

    Assuming a streamlit application app_1 under a default group named as
        streamlit_default is configured, this app can be created as:

        >>> from snowflake_ai.apps import StreamlitApp
        ... 
        ... # initialize application config for app_1
        >>> app = StreamlitApp("streamlit_default.app_1")
        ...
        ... # optionally the configs may be loaded from your custom dir
        ... # app = AppConfig("group_0.app_1", "custom dir", "config.toml")
    """

    _logger = logging.getLogger(__name__)
    _logger.addHandler(logging.StreamHandler(sys.stdout))

    K_PAGE = "page"
    K_PRE_PAGE = "previous_page"
    T_ST_APP = "streamlit"


    def __init__(
        self,
        app_key : str,
        config_dir : Optional[Union[str, None]] = None, 
        config_file : Optional[Union[str, None]] = None
    ):
        """
        Creat app specific AppConfig object.

        Args:
            app_key: A string representing the application; it can have 
                the format of <app_group>.<application> ('<', '>' not 
                included)
            config_dir: directory path string for custom config load
            config_file: file name string for custom config load
        """
        super().__init__(app_key, config_dir, config_file)
        self.logger = AppConfig._logger
        if StreamlitApp.T_ST_APP not in self.type:
            raise TypeError(
                "StreamlitApp.init(): Application type configuration Error."
            )
        self._setup_streamlit_config()
        self.pages : Dict[str, AppPage] = {}
        self.is_logged_in = False


    def _setup_streamlit_config(
        self, 
        conf_file: Optional[str] = None,
        file_content: Optional[str] = None
    ) -> str:
        """
        Setup streamlit app config file.

        Args:
            conf_file (str): config file name
            file_content (str): custom config file content
        
        Retrun:
            str: streamlit configuration file path
        """
        app_dir = os.path.join(self.root_path, self.script_home_dir)
        p = os.path.join(app_dir, ".streamlit/")
        if not exists(p):
            try:
                os.makedirs(p, exist_ok=True)
            except Exception as e:
                raise ValueError(
                    f"Streamlit._setup_streamlit_dir(): Error - {e}. "\
                    "Cannot create streamlit configuration directory"
                )

        if conf_file is None:
            conf_file = "config.toml"
        
        conf_path = os.path.join(p, conf_file)

        if not exists(conf_path):
            if file_content is None:
                file_content = self.get_streamlit_config()
            
            if file_content:
                with open(conf_path, "w") as f:
                    f.write(file_content)
        else:
            with open(conf_path, "r") as f:
                current_content = f.read()
            
            if file_content is None:
                file_content = self.get_streamlit_config()
            
            if current_content != file_content:
                if file_content:
                    with open(conf_path, "w") as f:
                        f.write(file_content)

        return conf_path


    def get_streamlit_config(self) :
        """
        Get streamlit app configuration from config file
        """
        # split app_key in form of group_key.app_key into two variables
        gk, k = AppConfig.split_group_key(self.app_key)
        config: Dict[str, Dict] = self.get_all_configs().get(
            AppConfig.K_ST_APPS, {}
        )
        self.logger.info(
            f"StreamlitApp.get_streamlit_config(): group_key =>{gk}; app_"\
            f"key => {k}; config => {config}"
        )
        if config:
            config = config.get(gk, {})
            if config:
                config = config.get(k, {})
        
        return toml.dumps(config)
    

    def add_page(self, page:AppPage):
        if self.pages:
            self.pages[page.page_id] = page
        

    def _add_page(
        self,
        page_id: str,
        title: Optional[str] = None,
        icon: Optional[str] = None,
        layout: Optional[str] = "wide",
        show_sidebar: Optional[bool] = True,
        func: Optional[Callable] = None,
    ):
        page = AppPage(
            page_id=page_id,
            title=title,
            icon=icon,
            layout=layout,
            show_sidebar=show_sidebar,
            func=func,
        )
        self.pages[page_id] = page


    def register_page(
        self,
        page_id: str,
        title: Optional[str] = None,
        icon: Optional[str] = None,
        layout: Optional[str] = "wide",
        show_sidebar: Optional[bool] = True,
    ) -> Callable:
        def decorator(func: Callable) -> Callable:
            @wraps(func)
            def wrapper(*args, **kwargs):
                return func(*args, **kwargs)

            self._add_page(
                page_id=page_id,
                title=title,
                icon=icon,
                layout=layout,
                show_sidebar=show_sidebar,
                func=wrapper,
            )
            return wrapper

        return decorator
    

    def set_pages_config(self):
        for _, pg in self.pages.items():
            st.set_page_config(
                page_title = pg.title,
                page_icon = pg.icon ,
                layout = pg.layout,
            )


    @staticmethod
    def get_current_page():
        query_params = st.experimental_get_query_params()
        return query_params.get(StreamlitApp.K_PAGE, ["/"])[0]
    

    def navigate_to(self, page_id: str):
        previous_page_id = self.get_current_page()
        st.experimental_set_query_params(
            previous_page = previous_page_id
        )
        st.experimental_set_query_params(page=page_id)
        self.run()


    def run(self):
        current_page_id = self.get_current_page()
        page: AppPage = self.pages.get(current_page_id)
        if page:
            page.func()
        else:
            st.error("Page not found!")


    def request_auth_code(
        self,
        oc:OAuthConnect
    ) -> str:
        oc = ConnectManager.create_default_oauth_connect(self)
        _, cc = AuthCodeConnect.generate_pkce_pair()
        params = {
            AuthCodeConnect.K_CODE_CHALLENGE: cc
        }
        return oc.authorize_request(params)


    def request_access_token(
        self,
        oc: OAuthConnect,
        ap: AppPage = None
    ) -> Dict :
        ok: bool = False
        params = oc.prepare_grant_request()
        auth_cd = params.get("auth_code")
        ctx = {}
        if auth_cd is not None and auth_cd:
            tok_res = oc.grant_request(params)
            print(f"DEBUG S0 ========> token result: {tok_res}")
            if tok_res:
                dc = ConnectManager.create_default_snow_connect(self)
                ctx = oc.decode_token(
                    tok_res, ["access_token", "refresh_token"]
                )
                print(f"DEBUG S1 => context => {ctx.items}")
                session = dc.create_session(ctx)
                if ap is not None:
                    ap.session = session
                    print(f"DEBUG S2 => session: true")

                tok: Dict = ctx.get("decoded_access_token")
                if tok:
                    fnm = tok.get("given_name", "")
                    lnm = tok.get("family_name", "")
                    st.sidebar.write(
                        f"Welcome, {fnm} {lnm}!"
                    )
                ok = True
        if not ok:
            st.error("Authorization failed.")

        return ctx


    def request_refresh_token(
        self,
        ctx: Dict,
        oc: OAuthConnect,
        ap: AppPage = None
    ) -> Dict :
        ok: bool = False
        params: Dict = oc.prepare_token_refresh(ctx) 
        refresh_tok = params.get("refresh_token")
        ctx = {}
        if  refresh_tok is not None and refresh_tok:
            tok_res = oc.refresh_token_request(params)
            print(f"DEBUG S1 ===========> refersh result: {tok_res}")
            if tok_res:
                dc = ConnectManager.create_default_snow_connect(self)
                ctx = oc.decode_token(
                    tok_res, ["access_token", "refresh_token"]
                )
                if ctx.get("access_token"):
                    print(f"DEBUG S2 ===========> context: true")
                session = dc.create_session(ctx)
                if ap is not None:
                    ap.session = session

                tok: Dict = ctx.get("decoded_access_token")
                if tok:
                    fnm = tok.get("given_name", "")
                    lnm = tok.get("family_name", "")
                    st.sidebar.write(
                        f"Welcome, {fnm} {lnm}!"
                    )
                ok = True
        if not ok:
            st.error("Authorization failed.")

        return ctx
    