import streamlit as st
import os
import datetime as DT
import pytz
import time
import json
import re
from typing import List
from transformers import AutoTokenizer
from gradio_client import Client

from dotenv import load_dotenv
load_dotenv()

useGpt4 = os.environ.get("USE_GPT_4") == "1"

if useGpt4:
    from openai import OpenAI
    client = OpenAI(api_key=os.environ.get("OPENAI_API_KEY"))
    MODEL = "gpt-4o-mini"
    MAX_CONTEXT = 128000
    tokenizer = AutoTokenizer.from_pretrained("Xenova/gpt-4o")
else:
    from groq import Groq
    client = Groq(
        api_key=os.environ.get("GROQ_API_KEY"),
    )
    MODEL = "llama-3.1-70b-versatile"
    MAX_CONTEXT = 8000
    tokenizer = AutoTokenizer.from_pretrained("Xenova/Meta-Llama-3.1-Tokenizer")


JSON_SEPARATOR = ">>>>"


def countTokens(text):
    text = str(text)
    tokens = tokenizer.encode(text, add_special_tokens=False)
    return len(tokens)


SYSTEM_MSG = f"""
You're an storytelling assistant who guides users through four phases of narrative development, helping them craft compelling personal or professional stories. The story created should be in simple language, yet evoke great emotions.
Ask one question at a time, give the options in a numbered and well formatted manner in different lines
If your response has number of options to choose from, only then append your final response with this exact keyword "{JSON_SEPARATOR}", and only after this, append with the JSON of options to choose from. The JSON should be of the format:
{{
    "options": [
        {{ "id": "1", "label": "Option 1"}},
        {{ "id": "2", "label": "Option 2"}}
    ]
}}
Do not write "Choose one of the options below:"
Keep options to less than 9.
Summarise options chosen so far in each step.

# Tier 1: Story Creation
You initiate the storytelling process through a series of engaging prompts:
Story Origin:
Asks users to choose between personal anecdotes or adapting a well-known story (creating a story database here of well-known stories to choose from).

Story Use Case:
Asks users to define the purpose of building a story (e.g., profile story, for social media content).

Story Time Frame:
Allows story selection from various life stages (childhood, mid-career, recent experiences).
Or Age-wise (below 8, 8-13, 13-15 and so on).

Story Focus:
Prompts users to select behaviours or leadership qualities to highlight in the story.
Provides a list of options based on common leadership traits:
(Generosity / Integrity / Loyalty / Devotion / Kindness / Sincerity / Self-control / Confidence / Persuasiveness / Ambition / Resourcefulness / Decisiveness / Faithfulness / Patience / Determination / Persistence / Fairness / Cooperation / Optimism / Proactive / Charisma / Ethics / Relentlessness / Authority / Enthusiasm / Boldness)

Story Type:
Prompts users to select the kind of story they want to tell:
Where we came from: A founding Story
Why we can't stay here: A case-for-change story
Where we're going: A vision story
How we're going to get there: A strategy story
Why I lead the way I do: Leadership philosophy story
Why you should want to work here: A rallying story
Personal stories: Who you are, what you do, how you do it, and who you do it for
What we believe: A story about values
Who we serve: A customer story
What we do for our customers: A sales story
How we're different: A marketing story

Guided Storytelling Framework:
You then lead users through a structured narrative development via the following prompts:
- Describe the day it happened
- What was the Call to Action / Invitation
- Describing the obstacles (up to three) in 4 lines
- Exploring emotions/fears experienced during the incident
- Recognize the helpers / any objects of help in the incident
- Detailing the resolution / Reaching the final goal
- Reflecting on personal growth or lessons learned (What did you do that changed your life forever?)

Now, show the story created so far, and ask for confirmation before proceeding to the next tier.

# Tier 2: Story Enhancement
After initial story creation, you offer congratulations on completing the first draft and gives 2 options:
Option 1 - Provides option for one-on-one sessions with expert storytelling coaches - the booking can be done that at https://calendly.com/
Options 2 - Provides further options for introducing users to more sophisticated narratives.

If Option 2 chosen, show these options with simple explanation and chose one.
You take the story and integrates it into different options of storytelling narrative structure:
The Story Hanger
The Story Spine
Hero's Journey
Beginning to End / Beginning to End
In Media Res (Start the story in the middle)
Nested Loops
The Cliffhanger

After taking user's preference, you show the final story and ask for confirmation before moving to the next tier.
Allow them to iterate over different narratives to see what fits best for them.

# Tier 3: Story Polishing
The final phase focuses on refining the narrative further:
You add suggestions to the story:
Impactful quotes/poems / similes/comparisons
Creative enhancements:
Some lines or descriptions for inspiration
Tips for maximising emotional resonance and memorability
By guiding users through these three tiers, you aim to cater to novice storytellers, offering a comprehensive platform for narrative skill development through its adaptive approach.
You end it with the final story and seeking any suggestions from the user to refine the story further.
Once the user confirms, you congratulate them with emojis on completing the story and provide the final story in a beatifully formatted manner.
Note that the final story should include twist, turns and events that make it really engaging and enjoyable to read.

"""

USER_ICON = "man.png"
AI_ICON = "Kommuneity.png"
IMAGE_LOADER = "ripple.svg"
TEXT_LOADER = "balls.svg"
START_MSG = "I want to create a story 😊"

st.set_page_config(
    page_title="Kommuneity Story Creator",
    page_icon=AI_ICON,
    # menu_items={"About": None}
)
ipAddress = st.context.headers.get("x-forwarded-for")


def __nowInIST() -> DT.datetime:
    return DT.datetime.now(pytz.timezone("Asia/Kolkata"))


def pprint(log: str):
    now = __nowInIST()
    now = now.strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{now}] [{ipAddress}] {log}")


pprint("\n")

st.markdown(
    """
    <style>
    @keyframes blinker {
        0% {
            opacity: 1;
        }
        50% {
            opacity: 0.2;
        }
        100% {
            opacity: 1;
        }
    }

    .blinking {
        animation: blinker 3s ease-out infinite;
    }

    .code {
        color: green;
        border-radius: 3px;
        padding: 2px 4px; /* Padding around the text */
        font-family: 'Courier New', Courier, monospace; /* Monospace font */
    }

    </style>
    """,
    unsafe_allow_html=True
)


def __isInvalidResponse(response: str):
    # new line followed by small case char
    if len(re.findall(r'\n[a-z]', response)) > 3:
        return True

    # lot of repeating words
    if len(re.findall(r'\b(\w+)(\s+\1){2,}\b', response)) > 1:
        return True

    # lots of paragraphs
    if len(re.findall(r'\n\n', response)) > 15:
        return True

    # json response without json separator
    if ('{\n    "options"' in response) and (JSON_SEPARATOR not in response):
        return True


def __matchingKeywordsCount(keywords: List[str], text: str):
    return sum([
        1 if keyword in text else 0
        for keyword in keywords
    ])


def __isStringNumber(s: str) -> bool:
    try:
        float(s)
        return True
    except ValueError:
        return False


def __getImagePromptDetails(prompt: str, response: str):
    regex = r'[^a-z0-9 \n\.\-]|((the) +)'

    cleanedResponse = re.sub(regex, '', response.lower())
    pprint(f"{cleanedResponse=}")

    cleanedPrompt = re.sub(regex, '', prompt.lower())
    pprint(f"{cleanedPrompt=}")

    if (
        __matchingKeywordsCount(
            ["adapt", "profile", "social media", "purpose", "use case"],
            cleanedResponse
        ) > 2
        and not __isStringNumber(prompt)
        and cleanedPrompt in cleanedResponse
        and "story so far" not in cleanedResponse
    ):
        return (
            f'''
                Subject: {prompt}.
                Style: Fantastical, in a storybook, surreal, bokeh
            ''',
            "Painting your character ..."
        )

        '''
        Mood: ethereal lighting that emphasizes the fantastical nature of the scene.

        storybook style

        4d model, unreal engine

        Alejandro Bursido

        vintage, nostalgic

        Dreamlike, Mystical, Fantastical, Charming
        '''

    if __matchingKeywordsCount(
        ["tier 2", "tier-2"],
        cleanedResponse
    ) > 0:
        possibleStoryEndIdx = [response.find("tier 2"), response.find("tier-2")]
        storyEndIdx = max(possibleStoryEndIdx)
        relevantResponse = response[:storyEndIdx]
        pprint(f"{relevantResponse=}")
        return (
            f"photo of a scene from this text: {relevantResponse}",
            "Imagining your scene (beta) ..."
        )

    return (None, None)


def __resetButtonState():
    st.session_state["buttonValue"] = ""


def __setStartMsg(msg):
    st.session_state.startMsg = msg


if "chatHistory" not in st.session_state:
    st.session_state.chatHistory = []

if "messages" not in st.session_state:
    st.session_state.messages = []

if "buttonValue" not in st.session_state:
    __resetButtonState()

if "startMsg" not in st.session_state:
    st.session_state.startMsg = ""


def __getMessages():
    def getContextSize():
        currContextSize = countTokens(SYSTEM_MSG) + countTokens(st.session_state.messages) + 100
        pprint(f"{currContextSize=}")
        return currContextSize

    while getContextSize() > MAX_CONTEXT:
        pprint("Context size exceeded, removing first message")
        st.session_state.messages.pop(0)

    return st.session_state.messages


def predict():
    messagesFormatted = [{"role": "system", "content": SYSTEM_MSG}]
    messagesFormatted.extend(__getMessages())
    contextSize = countTokens(messagesFormatted)
    pprint(f"{contextSize=} | {MODEL}")

    response = client.chat.completions.create(
        model=MODEL,
        messages=messagesFormatted,
        temperature=0.8,
        max_tokens=4000,
        stream=True
    )

    chunkCount = 0
    for chunk in response:
        chunkContent = chunk.choices[0].delta.content
        if chunkContent:
            chunkCount += 1
            yield chunkContent


def generateImage(prompt: str):
    pprint(f"imagePrompt={prompt}")
    fluxClient = Client("black-forest-labs/FLUX.1-schnell")
    result = fluxClient.predict(
            prompt=prompt,
            seed=0,
            randomize_seed=True,
            width=1024,
            height=768,
            num_inference_steps=4,
            api_name="/infer"
    )
    pprint(f"imageResult={result}")
    return result


st.title("Kommuneity Story Creator 📖")
if not (st.session_state["buttonValue"] or st.session_state["startMsg"]):
    st.button(START_MSG, on_click=lambda: __setStartMsg(START_MSG))

for chat in st.session_state.chatHistory:
    role = chat["role"]
    content = chat["content"]
    imagePath = chat.get("image")
    avatar = AI_ICON if role == "assistant" else USER_ICON
    with st.chat_message(role, avatar=avatar):
        st.markdown(content)
        if imagePath:
            st.image(imagePath)

if prompt := (st.chat_input() or st.session_state["buttonValue"] or st.session_state["startMsg"]):
    __resetButtonState()
    __setStartMsg("")

    with st.chat_message("user", avatar=USER_ICON):
        st.markdown(prompt)
    pprint(f"{prompt=}")
    st.session_state.messages.append({"role": "user", "content": prompt})
    st.session_state.chatHistory.append({"role": "user", "content": prompt })

    with st.chat_message("assistant", avatar=AI_ICON):
        responseContainer = st.empty()

        def __printAndGetResponse():
            response = ""
            # responseContainer.markdown(".....")
            responseContainer.image(TEXT_LOADER)
            responseGenerator = predict()

            for chunk in responseGenerator:
                response += chunk
                if __isInvalidResponse(response):
                    pprint(f"{response=}")
                    return

                if JSON_SEPARATOR not in response:
                    responseContainer.markdown(response)

            return response

        response = __printAndGetResponse()
        while not response:
            pprint("Empty response. Retrying..")
            time.sleep(0.5)
            response = __printAndGetResponse()

        pprint(f"{response=}")

        def selectButton(optionLabel):
            st.session_state["buttonValue"] = optionLabel
            pprint(f"Selected: {optionLabel}")

        responseParts = response.split(JSON_SEPARATOR)

        jsonStr = None
        if len(responseParts) > 1:
            [response, jsonStr] = responseParts

        imagePath = None
        imageContainer = st.empty()
        try:
            (imagePrompt, loaderText) = __getImagePromptDetails(prompt, response)
            if imagePrompt:
                imgContainer = imageContainer.container()
                imgContainer.write(
                    f"""
                    <div class='blinking code'>
                    {loaderText}
                    </div>
                    """,
                    unsafe_allow_html=True
                )
                # imgContainer.markdown(f"`{loaderText}`")
                imgContainer.image(IMAGE_LOADER)
                (imagePath, seed) = generateImage(imagePrompt)
                imageContainer.image(imagePath)
        except Exception as e:
            pprint(e)
            imageContainer.empty()

        if jsonStr:
            try:
                json.loads(jsonStr)
                jsonObj = json.loads(jsonStr)
                options = jsonObj["options"]

                for option in options:
                    st.button(
                        option["label"],
                        key=option["id"],
                        on_click=lambda label=option["label"]: selectButton(label)
                    )
                # st.code(jsonStr, language="json")
            except Exception as e:
                pprint(e)

        st.session_state.messages.append({
            "role": "assistant",
            "content": response,
        })
        st.session_state.chatHistory.append({
            "role": "assistant",
            "content": response,
            "image": imagePath,
        })