from datetime import date
import os
import re
from time import sleep

from bardapi import BardCookies
from dotenv import load_dotenv
import google.generativeai as Gemini
from openai import OpenAI
from openai.types import FileObject as OpenAIFile
from openai.types.beta import Assistant as OpenAIAssistant
from openai.types.beta import Thread as OpenAIThread
import streamlit as st


load_dotenv()


def get_bard_client() -> BardCookies:
    return BardCookies(cookie_dict={
        "__Secure-1PAPISID": os.environ.get("BARD_Secure-1PAPISID"),
        "__Secure-1PSID": os.environ.get("BARD_Secure-1PSID"),
        "__Secure-1PSIDCC": os.environ.get("BARD_Secure-1PSIDCC"),
        "__Secure-1PSIDTS": os.environ.get("BARD_Secure-1PSIDTS"),
    })


def ask_bard(status_writer, bard_client: BardCookies, content: str) -> str:
    """Ask Bard a question."""
    status_writer.write("Asking Bard...")
    response = bard_client.get_answer(content)
    return response["content"]


def get_gemini_client() -> Gemini.GenerativeModel:
    Gemini.configure(api_key=os.environ.get("GEMINI_API_KEY"))
    return Gemini.GenerativeModel('gemini-pro')


def ask_gemini(status_writer, gemini_client: Gemini.GenerativeModel, content: str) -> None:
    """Ask Gemini a question."""
    messages = [
        {
            'role': 'user',
            'parts': content,
        }
    ]

    status_writer.write("Asking Gemini...")
    response = ""
    partial_response = gemini_client.generate_content(
        messages,
        stream=True
    )
    for chunk in partial_response:
        response += chunk.candidates[0].content.parts[0].text
    return response


def get_sf_client():
    return st.connection("snowflake")


def run_sf_callable(sf_client, sql):
    message["results"] = sf_client.query(sql)
    return message["results"]


openai_client = OpenAI(api_key=os.environ.get("OPENAPI_API_KEY"))

OPENAPI_SLEEP: int = 10
OPENAPI_MODELS = [
    "gpt-4-0125-preview",  # latest 128,000 tokens Up to Apr 2023
    "gpt-4-1106-preview",  # 128,000 tokens Up to Apr 2023
    "gpt-3.5-turbo-0125",  # 16,385 tokens Up to Sep 2021
]

ASSISTANT_NAMES = [
    "Skittles",  # The AI Snowflake SQL Expert
    "Charles",  # Just fetches what it can
]
ASSISTANT_MODEL = OPENAPI_MODELS[0]
ASSISTANT_NAME = ASSISTANT_NAMES[0]
OPENAPI_INSTRUCTIONS = {
    "Skittles": f"""Today is {date.today().strftime("%B %d, %Y")}.
    
You are Skittles, an AI Snowflake SQL copilot. Your responses are always in the character of Skittles.

You never hallucinate about database schemas or models, and will ask for clarification.
Users will sometimes ask you questions about the database data or to write a query, and in response,
you will return valid Snowflake SQL, enclosed with sql markdown, based on the database schemas in your Knowledgebase.
When providing SQL, try to use columns which represent universal identifiers instead of internal identifiers, unless
specified by the user. When providing SQL that describes audio/music/video streams (aka tracks), include the track name, isrc, performers and label name.

Should a user ask you a question about recent trends, or anything you might not have training data for, respond with "Ask Bard!".
""",
    "Charles": "",
}

ASSISTANT_TOOLS = [
#     {"type": "code_interpreter"},
#     {"type": "function"},
    {"type": "retrieval"},
]
ASSISTANT_FILES = ["kb.md"]


def _manage_assistant_kb(
    openai_client: OpenAI,
    assistant: OpenAIAssistant
) -> OpenAIAssistant:
    """Ensure assistant has access to ASSISTANT_FILES"""
    assistant_file_lookup = {}
    for file_id in assistant.file_ids:
        file = openai_client.files.retrieve(file_id=file_id)
        assistant_file_lookup[file.filename] = file

    for filename in ASSISTANT_FILES:
        if not filename in assistant_file_lookup:
            # Get or create the file
            file: OpenAIFile = None
            file = openai_client.files.create(
                file=open(filename, "rb"),
                purpose='assistants'
            )
            # Give the file to the assistant
            openai_client.beta.assistants.files.create(
                assistant_id=assistant.id,
                file_id=file.id,
            )

    assistant = openai_client.beta.assistants.retrieve(assistant_id=assistant.id)
    return assistant


def get_assistant(openai_client: OpenAI, status_writer) -> OpenAIAssistant:
    """Create/Update an Assistant with tools, knowledge, etc."""
    assistants = openai_client.beta.assistants.list()
    assistant: OpenAIAssistant = None

    for candidate in assistants:
        if candidate.name == ASSISTANT_NAME:
            assistant = candidate
            status_writer.write(f"Updating assistant {assistant.id}...")
    if not assistant:
        # Create the assistant
        assistant = openai_client.beta.assistants.create(
            model=ASSISTANT_MODEL,
        )
        status_writer.write(f"Created assistant {assistant.id}...")

    # Make sure the assistant is consistent, whether just created or existing
    assistant = openai_client.beta.assistants.update(
        assistant_id=assistant.id,
        model=ASSISTANT_MODEL,
        name=ASSISTANT_NAME,
        instructions=OPENAPI_INSTRUCTIONS[ASSISTANT_NAME],
        tools=ASSISTANT_TOOLS,
    )

    if ASSISTANT_NAME == "Skittles":
        assistant = _manage_assistant_kb(openai_client, assistant)

    status_writer.write(f"Meet {assistant.name} ({assistant.id}), running on {assistant.model}")
    status_writer.write(f"KB Files added: {', '.join([file_id for file_id in assistant.file_ids])}")

    return assistant


def ask_assistant(
    status_writer,
    assistant: OpenAIAssistant,
    thread: OpenAIThread,
    content: str
) -> str:
    # Put the message on the thread
    openai_client.beta.threads.messages.create(
        thread_id=thread.id,
        role="user",
        content=content,
    )
    # Associate the thread and assistant together and ask
    run = openai_client.beta.threads.runs.create(
        thread_id=thread.id,
        assistant_id=assistant.id,
    )
    status_writer.write(f"Waiting for the assistant... The run id is {run.id}, status: {run.status}")

    # TODO Handle other statuses than "completed":
    # queued, in_progress, requires_action, cancelling, cancelled, failed, completed, or expired.
    while(run.status != "completed"):
        if run.status == "failed":
            status_writer.write("Start using the debugger")
            import pdb; pdb.set_trace()
            break
        sleep(OPENAPI_SLEEP)
        run = openai_client.beta.threads.runs.retrieve(
            thread_id=thread.id,
            run_id=run.id,
        )
        status_writer.write(f"Assistant status: {run.status}")

    messages = openai_client.beta.threads.messages.list(thread_id=thread.id)
    return messages.data[0].content[0].model_dump()["text"]["value"]

# Initialize everything!
st.title(f"{ASSISTANT_NAME} is here to help!")
st_status_container = st.status("Things are happening though you can't see it...")
if "messages" not in st.session_state:
    st.session_state.messages = [{"role": "system", "content": f"{ASSISTANT_NAME} is here to help you with all Analytics questions."}]

st_status_container.write("Initializing snowflake connection...")
st.session_state.sf_client = get_sf_client()
st_status_container.write("Initializing bard client...")
st.session_state.bard_client = get_bard_client()
st_status_container.write("Initializing gemini client...")
st.session_state.gemini_client = get_gemini_client()
st_status_container.write("Initializing openai assistant...")
st.session_state.assistant = get_assistant(openai_client=openai_client, status_writer=st_status_container)
st_status_container.write("Initializing openai thread...")
st.session_state.thread = openai_client.beta.threads.create()
st_status_container.write(f"The Thread Id is {st.session_state.thread.id}")

st.session_state.is_waiting_for_user = True

# Prompt for user input and save
if prompt := st.chat_input():
    st.session_state.is_waiting_for_user = False
    st.session_state.messages.append({"role": "user", "content": prompt})

# Display the chat messages in streamlit
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.write(message["content"])
        if "results" in message:
            st.dataframe(message["results"])


# If we are not waiting for the user, generate a response for the user's prompt
if not st.session_state.is_waiting_for_user:  # st.session_state.messages[-1]["role"] == "user":
    prompt = st.session_state.messages[-1]["content"]
    # Take the message and send first to Assistant
    response = ask_assistant(
        status_writer=st_status_container,
        assistant=st.session_state.assistant,
        thread=st.session_state.thread,
        content=prompt,
    )

    # Check response for "Ask Bard!"
    bard_match = re.search(r"Ask Bard!", response, re.DOTALL)
    if bard_match:
        with st.chat_message("assistant"):
            # bard_says = ask_gemini(st_status_container, st.session_state.gemini_client, prompt)
            bard_says = ask_bard(
                status_writer=st_status_container,
                bard_client=st.session_state.bard_client,
                content=prompt)
            st.write(bard_says)
            st.session_state.messages.append({"role": "assistant", "content": bard_says})
            st.session_state.is_waiting_for_user = True
    else:
        with st.chat_message("assistant"):
            st.session_state.messages.append({"role": "assistant", "content": response})
            st.write(response)
        # Parse the response for a SQL query and execute if available
        sql_match = re.search(r"```sql\n(.*)\n```", response, re.DOTALL)
        if sql_match:
            sql = sql_match.group(1)
            result = run_sf_callable(sf_client=st.session_state.sf_client, sql=sql)
            st.dataframe(result)
            st.session_state.messages.append({"role": "assistant", "content": result})
            st.session_state.is_waiting_for_user = True
        else:
            st.session_state.is_waiting_for_user = True