import glob
import gradio as gr
from typing import Any
from dotenv import load_dotenv
from griptape.structures import Agent
from griptape.tasks import PromptTask
from griptape.drivers import (
    LocalConversationMemoryDriver,
    GriptapeCloudStructureRunDriver,
    LocalFileManagerDriver,
    LocalStructureRunDriver,
)
from griptape.memory.structure import ConversationMemory
from griptape.tools import StructureRunClient, FileManager
from griptape.rules import Rule, Ruleset
from griptape.config import AnthropicStructureConfig
import time
import os
import re


# Load environment variables
load_dotenv()


# Create an agent that will create a prompt that can be used as input for the query agent from the Griptape Cloud.


# Function that logs user history - adds to history parameter of Gradio
# TODO: Figure out the exact use of this function
def user(user_message, history):
    history.append([user_message, None])
    return ("", history)


# Function that logs bot history - adds to the history parameter of Gradio
# TODO: Figure out the exact use of this function
def bot(history):
    response = send_message(history[-1][0])
    history[-1][1] = ""
    for character in response:
        history[-1][1] += character

        time.sleep(0.005)

        yield history


def create_prompt_task(session_id: str, message: str) -> PromptTask:
    return PromptTask(
        f"""
            Re-structure the values from the user's questions: '{message}' and the input value from the conversation memory '{session_id}.json' to fit the following format. Leave out attributes that aren't important to the user: 
                years experience: <x>
                location: <x>
                role: <x>
                skills: <x>
                expected salary: <x>
                availability: <x>
                past companies: <x>
                past projects: <x>
                show reel details: <x>
            """,
    )


def build_talk_agent(session_id: str, message: str) -> Agent:
    ruleset = Ruleset(
        name="Local Gradio Agent",
        rules=[
            Rule(
                value="You are responsible for structuring a user's questions into a specific format for a query."
            ),
            Rule(
                value="You ask the user follow-up questions to fill in missing information for the format you are trying to fit."
            ),
            Rule(
                value="If the user has no preference for a specific attribute, then you can remove it from the query."
            ),
            Rule(
                value="Only return the current query structure and any questions to fill in missing information."
            ),
            Rule(
                value="Don't return names or usernames, just return the initials of the responses."
            ),
        ],
    )
    file_manager_tool = FileManager(
        name="FileManager",
        file_manager_driver=LocalFileManagerDriver(),
        off_prompt=False,
    )

    return Agent(
        config=AnthropicStructureConfig(),
        conversation_memory=ConversationMemory(
            driver=LocalConversationMemoryDriver(file_path=f"{session_id}.json")
        ),
        tools=[file_manager_tool],
        tasks=[create_prompt_task(session_id, message)],
        rulesets=[ruleset],
    )


# Creates an agent for each run
# The agent uses local memory, which it differentiates between by session_hash.
def build_agent(session_id: str, message: str) -> Agent:

    ruleset = Ruleset(
        name="Local Gradio Agent",
        rules=[
            Rule(
                value="You are responsible for structuring a user's questions into a specific format for a query and then querying."
            ),
            Rule(
                value="Only return the result of the query, do not provide additional commentary."
            ),
            Rule(value="Only perform one task at a time."),
            Rule(
                value="Do not perform the query unless the user has said 'Done' with formulating."
            ),
            Rule(
                value="Only perform the query with the proper query structure as one string argument."
            ),
            Rule(
                value="If you reformulate the query, then you must ask the user if they are 'Done' again."
            ),
            Rule(
                value="If the user says they want to start over, then you must delete the conversation memory file."
            ),
        ],
    )

    print("Base URL", os.environ.get("BASE_URL", "https://cloud.griptape.ai"))

    query_client = StructureRunClient(
        name="QueryResumeSearcher",
        description="Use it to search for a candidate with the query.",
        driver=GriptapeCloudStructureRunDriver(
            # base_url=os.environ.get("BASE_URL","https://cloud.griptape.ai"),
            structure_id=os.getenv("GT_STRUCTURE_ID"),
            api_key=os.getenv("GT_CLOUD_API_KEY"),
            structure_run_wait_time_interval=5,
            structure_run_max_wait_time_attempts=30,
        ),
    )

    talk_client = StructureRunClient(
        name="FormulateQueryFromUser",
        description="Used to formulate a query from the user's input.",
        driver=LocalStructureRunDriver(
            structure_factory_fn=lambda: build_talk_agent(session_id, message),
        ),
    )

    return Agent(
        config=AnthropicStructureConfig(),
        conversation_memory=ConversationMemory(
            driver=LocalConversationMemoryDriver(file_path=f"{session_id}.json")
        ),
        tools=[talk_client, query_client],
        rulesets=[ruleset],
    )


def delete_json(session_id: str) -> None:
    for file in glob.glob(f"{session_id}.json"):
        os.remove(file)


def send_message(message: str, history, request: gr.Request) -> Any:
    if request:
        session_hash = request.session_hash
        agent = build_agent(session_hash, message)
    response = agent.run(message)
    # if re.search(r'\bdone[.,!?]?\b', message, re.IGNORECASE):
    # delete_json(session_hash)
    return response.output.value


demo = gr.ChatInterface(fn=send_message)
demo.launch()