import streamlit as st
from melvin.tecton_gen_ai.api import Agent
from utils import _log, load_object, store_object
from typing import Callable
from editors.feature_view_editor_ui import display_editor, FeatureEditorState

def _load_fv_editor_state(session_object_key: str) -> FeatureEditorState:
    try:
        state = load_object(session_object_key)
        state.session_object_key = session_object_key
        return state
    except Exception as e:
        _log(f"Error loading editor state: {e}")
        return FeatureEditorState(result_apply=False, result_cancel=False, session_object_key=session_object_key)

def _complete_session(editor_state: FeatureEditorState):
    st.session_state[editor_state.session_object_key] = True
    store_object(editor_state, "editor", editor_state.session_object_key)

def _create_fv_from_state(editor_state: FeatureEditorState, simulate_user_prompt_fn: Callable):
    submit_user_prompt = f"Create the following feature views for me. Create a new data source it doesn't exist. Create a new entity if it doesn't exist. Make up the FV name: {editor_state.to_dict()}"
    simulate_user_prompt_fn(submit_user_prompt)

def _on_apply(editor_state: FeatureEditorState, simulate_user_prompt_fn: Callable):
    editor_state.result_apply = True
    _complete_session(editor_state)
    _create_fv_from_state(editor_state, simulate_user_prompt_fn)

def _on_cancel(editor_state: FeatureEditorState):
    editor_state.result_cancel = True
    _complete_session(editor_state)


def _display_completed_session(editor_state: FeatureEditorState):
    st.code(f"Session completed: {editor_state.session_object_key}. Result: {editor_state.result_apply} {editor_state.result_cancel}")

def display_fv_editor(session_object_key: str, simulate_user_prompt_fn: Callable):
    editor_state = None

    # Load the editor state from the streamlit session state or from file if it was just created by the agent
    if f"{session_object_key}_ui_state" in st.session_state:
        editor_state = st.session_state[f"{session_object_key}_ui_state"]
    else:
        editor_state = _load_fv_editor_state(session_object_key)
        st.session_state[f"{session_object_key}_ui_state"] = editor_state

    if session_object_key in st.session_state:
        _display_completed_session(editor_state)
    else:
        x, y = st.columns(2)
        with x:
            st.code(session_object_key)
            display_editor(editor_state, key_prefix=session_object_key, on_ship_it=lambda: _on_apply(editor_state, simulate_user_prompt_fn))


def sys_prompt() -> str:
    return """
You are going to generate a UI-based feature view editor.

You will call one and only one editor tool based on the user's intent to output the feature view editor. If there's no clear user intent, make one up.

If not successful, you should return an error message starting with `Error: `
If successful, return a editor state representation which is a session object key starting with `so_editor_`,
and surround it with triple backticks, language is `tecton_editor`

The editor state must be generated by one of the tools. don't manipulate the `so_editor_` id.

Example of the editor state representation:

```tecton_editor
so_editor_sdgg
```
"""

def _generate_aggregation_feature_view_editor_state() -> str:
    """
    Generate aggregation feature view editor

        Returns:

            str: the editor state representation session object key, starting with `so_editor_`
    """

    editor_state = FeatureEditorState(
        result_apply=False,
        result_cancel=False,
    )
    return store_object(editor_state, "editor")


def make_fv_editor_agent() -> Agent:
    return Agent(
        name="editor_agent",
        description="""An agent that generates a graphical Feature View editor state

Returns:

    if not successful, it returns an error message starting with `Error: `
    if successful, it returns a Feature View editor state quoted with triple backticks, language is `tecton_editor`. You must show the editor state quoted in triple backticks to the user.

Note:

    - The output of this agent should never be used to call this agent again
    - For any follow up question, the input data needs to be regenerated
    - Do not try to interpret the content in a tecton_editor code block, use the exact code block in the final output.
""",
        prompt=sys_prompt,
        tools=[
            _generate_aggregation_feature_view_editor_state
        ],
    )
