import json
import time
from typing import Optional

import streamlit as st
from sqlalchemy import select

from synthetic_agents.database.config import get_db
from synthetic_agents.database.entity.agent_type import AgentType
from synthetic_agents.database.entity.prompt_template import PromptTemplate
from synthetic_agents.model.memory import Memory
from synthetic_agents.webapp.common.constants import SECONDS_BEFORE_PAGE_REFRESH


def display_agent_info(attributes: Optional[dict[str, str]], life_facts: Optional[list[Memory]]):
    """
    Displays agent's attributes and life facts on the screen.

    :param attributes: attributes.
    :param life_facts: life facts.
    """
    if attributes is not None:
        st.write("**Agent Attributes**")
        st.json(attributes)
    else:
        st.write("**Agent Attributes**")
        st.write("*:red[No attributes.]*")

    if life_facts is not None:
        life_facts_as_string = "\n".join(
            [
                json.dumps(m.to_simplified_json(["creation_timestamp", "content"]))
                for m in life_facts
            ]
        )
        st.write("**Life Facts**")
        st.code(life_facts_as_string, line_numbers=True)
    else:
        st.write("**Life Facts**")
        st.write("*:red[No life facts.]*")


@st.cache_data
def read_available_agent_types():
    """
    Reads agent types from the database.
    :return: types of agents.
    """
    db = next(get_db())
    records = db.execute(select(AgentType).order_by(AgentType.description)).all()
    agent_types = [r[0] for r in records]
    db.close()

    return agent_types


@st.cache_data
def read_available_prompt_templates(
    agent_type: str, application_type: str
) -> list[PromptTemplate]:
    """
    Reads prompt templates for the agent and application types.

    :param agent_type: type of the agent for which the template was developed.
    :param application_type: type of the application for which the template was developed.
    :return: a list of prompt templates.
    """
    db = next(get_db())
    records = db.execute(
        select(PromptTemplate)
        .where(
            PromptTemplate.agent_type == agent_type,
            PromptTemplate.application_type == application_type,
        )
        .order_by(PromptTemplate.version)
    ).all()
    prompt_templates = [r[0] for r in records]
    db.close()

    return prompt_templates


def wait_and_refresh(num_seconds: int = SECONDS_BEFORE_PAGE_REFRESH):
    """
    Waits a few seconds and refreshes the page.

    :param num_seconds: number of seconds to wait before refreshing. Defaults to
        SECONDS_BEFORE_PAGE_REFRESH.
    """
    time.sleep(num_seconds)
    st.rerun()
