controller.py

from util.state import StateHandler

import json
import pandas as pd
from requests import post, get
import streamlit as st
from st_aggrid import AgGrid, GridUpdateMode
from st_aggrid.grid_options_builder import GridOptionsBuilder


def request_handler(state: StateHandler) -> dict:
    """Manages the request and response from the service accounting for the type of request from the user input, as updated in the program StateHander.

    Args:
        state (StateHandler): The state of the dashboard.
    
    Returns:
        dict: The response from the service.
    """

    if not st.session_state.request_submitted:
        st.stop()

    st.write("___")
    st.subheader("② Service repsonses")

    match st.session_state.request_type:
        case "any":
            payloads = _get_payloads_by_count(
                state.properties.service_url,
                st.session_state.payload_count,
            )
        case "ids":
            payloads = _get_payloads(
                state.properties.service_url,
                st.session_state.request_ids,
            )
        case _:
            print("Invalid request type")
            st.stop()

    _save_payloads(payloads)

    for payload in payloads:
        js, button = st.columns([0.75, 0.25])
        request_id = payload.get("requestId")
        with js:
            with st.expander(f"{request_id}.json"):
                st.json(payload, expanded=True)

        with button:
            st.download_button(
                label="Download JSON",
                data=json.dumps(payload, indent=4),
                file_name=f"{request_id}.json",
                key=f"download_{request_id}",
                mime="json",
            )
    render_dataframe(payloads)


def render_dataframe(payloads: list[dict]):
    """Renders the payloads as a dataframe in the dashboard.
    
    Args:
        payloads (list[dict]): The payloads to render.
    """

    st.write("")
    table = pd.json_normalize(payloads)
    gd = GridOptionsBuilder.from_dataframe(table)
    gd.configure_pagination(enabled=True)
    gd.configure_selection(selection_mode="multiple")

    AgGrid(
        table,
        gridOptions=gd.build(),
        update_mode=GridUpdateMode.SELECTION_CHANGED,
        theme="material",
    )

    st.write("___")


def _get_payloads(url: str, request_ids: list) -> dict | None:
    """Sends a POST request to the service and handles the response.

    The POST request includes a list of requestIds which corresponed to a specific payload requested from the service.
    
    Args:
        url (str): The url of the service to request from.
        request_ids (list): The request ids to request from the service.

    Returns:
        dict | None: The response from the service.
    """
    
    response = post(
        url,
        json={"requests": request_ids},
        verify=False,
    )
    print(f"response: {response}")
    if response.status_code == 200:
        print(f"response.text: {response.text}")
        return json.loads(response.text).get("payloads")
    else:
        if response.status_code == 403:
            print(f"Unauthorised access to {url}")
        else:
            print(f"Invalid request for {url}")
        print(f"    -> {response.text}")
        return None
    

def _get_payloads_by_count(url: str, count: int) -> dict | None:
    """Sends a GET request to the service and handles the response.

    The GET request includes a count in the URL which indicates the number of payloads to return from the service.
    
    Args:
        url (str): The url of the service to request from.
        count (int): The number of payloads to request from the service.
    """
    
    response = get(
        f"{url}/{count}",
        verify=False,
    )
    print(f"response: {response}")
    if response.status_code == 200:
        print(f"response.text: {response.text}")
        return json.loads(response.text).get("payloads")
    else:
        if response.status_code == 403:
            print(f"Unauthorised access to {url}")
        else:
            print(f"Invalid request for {url}")
        print(f"    -> {response.text}")
        return None


def _save_payloads(payloads: list) -> None:
    """Saves the payloads to local storage.
    
    Args:
        payloads (list): The payloads to save.
    """

    for payload in payloads:
        request_id = payload.get("requestId")

        with open(f"resources/payloads/{request_id}.json", "w") as f:
            json.dump(payload, f, indent=4)