from dotenv import load_dotenv
import gradio as gr
import requests
from retry import retry
from requests.exceptions import ConnectTimeout, ReadTimeout
import os
from uuid import uuid4
from typing import Optional, Sequence, Tuple


"""gradio app for Chat Interface for DataStax Langflow calls"""


load_dotenv()


BASE_API_URL = "https://api.langflow.astra.datastax.com"
LANGFLOW_ID = "e99b525a-84dd-4cf2-bac6-8e5ebc6a7d17"
FLOW_ID = "bfbf7237-50dd-40dd-baba-cc680288d788"
DATASTAX_TOKEN = os.getenv('DATASTAX_TOKEN')

ENDPOINT = 'ai_traveller'

TWEAKS = {
    "Memory-fCuCs": {
        "n_messages": 100,
        "order": "Ascending",
        "sender": "Machine and User",
        "sender_name": "",
        "template": "{sender_name}: {text}"
    },
    "ChatInput-PN5E4": {},
    "Prompt-vGvVG": {},
    "GoogleGenerativeAIModel-6n0Ft": {},
    "ChatOutput-HjfF1": {}
}


@retry((ConnectTimeout, ReadTimeout), tries=3, delay=2)
def call_ai(
    message: str,
    history: Sequence[Tuple[str, str]],
    local_storage: str,
    endpoint: str = ENDPOINT,
    output_type: str = "chat",
    input_type: str = "chat",
    tweaks: Optional[dict] = None,
    application_token: Optional[str] = DATASTAX_TOKEN
) -> dict:
    """
    Run a flow with a given message and optional tweaks.

    :param message: The message to send to the flow
    :param endpoint: The ID or the endpoint name of the flow
    :param tweaks: Optional tweaks to customize the flow
    :return: The JSON response from the flow
    """
    del history
    api_url = f"{BASE_API_URL}/lf/{LANGFLOW_ID}/api/v1/run/{endpoint}"

    payload = {
        'session_id': local_storage[0],
        "input_value": message,
        "output_type": output_type,
        "input_type": input_type,
    }
    headers = None
    if tweaks:
        payload["tweaks"] = tweaks
    if application_token:
        headers = {
            "Authorization": "Bearer " + application_token,
            "Content-Type": "application/json"
        }

    response = requests.post(
        api_url,
        json=payload,
        headers=headers,
        timeout=120,
    )
    print(response, response.content)
    response = response.json()
    resp_message = ''
    try:
        for resp in response['outputs']:
            for _resp in resp['outputs']:
                for message in _resp['messages']:
                    resp_message = message['message']
    except KeyError as e:
        print(e)
        print(response)

    return resp_message


def generate_session_id():
    """Generate a unique session ID."""
    return str(uuid4())


with gr.Blocks() as demo:
    with gr.Column():
        gr.Image(
            "./aem_logo.jpeg",
            width=300,
            label="AEM",
            show_download_button=False,
            show_fullscreen_button=False,
            show_label=False,
            interactive=False,
            show_share_button=False,
        )

    # session = gr.State(["", ], time_to_live=24 * 60*60)
    session = gr.BrowserState(None, storage_key="session_id")
    show_session = gr.Textbox(
        visible=False,
        interactive=False,
        label="Session ID",
    )
    with gr.Accordion(
        'Chat With Agent',
    ) as chat:
        chat_interface = gr.ChatInterface(
            call_ai,
            type='messages',
            title="AI Beauty Partner",
            additional_inputs=[session],
            autofocus=True,
            fill_height=True,
        )

    def initialize_session(existing_id):
        """
        Runs ONCE per browser session on page load.
        Checks if an ID already exists in BrowserState (local storage).
        If not, generates a new one.
        Returns the ID to populate the BrowserState and other components.
        """
        if existing_id:
            print(f"Found existing ID in BrowserState: {existing_id}")
            return existing_id, existing_id
        else:
            new_id = generate_session_id()
            print(f"No existing ID found, returning NEW ID: {new_id}")
            return new_id, new_id

    demo.load(  # pylint: disable=no-member
        initialize_session,
        [session],
        outputs=[session, show_session],
    )
    demo.launch(share=False, debug=True,)