import random
import time
from dataclasses import asdict, dataclass
from typing import Callable, Literal

import mesop.labs as mel
import mesop as me
import requests
import json
from dataclasses import dataclass, field

from mongo_db_writer_main import log_prompt_and_response

Role = Literal["user", "bot"]
_INTRO_TEXT = """
# Mesop Markdown Editor Example

This example shows how to make a simple markdown editor.
""".strip()


@dataclass(kw_only=True)
class Note:
  """Content of note."""

  content: str = ""


_APP_TITLE = "# YuvaRaj's LLM-Playground"
_BOT_AVATAR_LETTER = "UV"
_EMPTY_CHAT_MESSAGE = "Get started with an example"
_EXAMPLE_USER_QUERIES = (
    "Examples of SQL Commands?",
    "What a store proc? Provide an examples",
    "What is Java Class, provide few examples?",
    "What is quantum computing, and why is it important?"
)
_CHAT_MAX_WIDTH = "800px"
_MOBILE_BREAKPOINT = 640


@dataclass(kw_only=True)
class ChatMessage:
    """Chat message metadata."""

    role: Role = "user"
    content: str = "You're ai expert, guided by yuvaraj"
    edited: bool = False
    # 1 is positive
    # -1 is negative
    # 0 is no rating
    rating: int = 0

@me.stateclass
class State:
    input: str
    output: list[ChatMessage]
    in_progress: bool
    sidebar_expanded: bool = False
    # Need to use dict instead of ChatMessage due to serialization bug.
    # See: https://github.com/google/mesop/issues/659
    history: list[list[dict]]
    notes: list[Note] = field(default_factory=lambda: [Note(content=_INTRO_TEXT)])
    selected_note_index: int = 0
    selected_note_content: str = _INTRO_TEXT
    show_preview: bool = True




def call_ollama_api(input_text: str, history: list[ChatMessage]) -> str:
    """
    Sends the user input and chat history to the Ollama endpoint and retrieves the response.
    """
    OLLAMA_URL = "http://96.228.37.139:4002/api/generate"
    MODEL_NAME = "mistral-small"

    # Format the chat history as part of the prompt
    formatted_history = "\n".join([f"{msg.role}: {msg.content}" for msg in history])
    prompt = f"{formatted_history}\nuser: give me response is well formatted  {input_text}"

    # Prepare the payload
    payload = {
        "model": MODEL_NAME,
        "prompt": prompt,  # Include chat history in the prompt
        "stream": False,
    }

    try:
        response = requests.post(OLLAMA_URL, json=payload)
        response.raise_for_status()

        # Parse the JSON response
        response_data = response.json()
        
        log_reponse_data=response_data.get("response")

        # Access the 'response' field from the parsed JSON data
        if "response" in response_data:
              # Call the external logging function
            log_prompt_and_response(prompt, log_reponse_data)
            return response_data.get("response")
        
      

        else:
            return "Error: No response field found in Ollama API response"

    except requests.exceptions.RequestException as e:
        return f"Error connecting to Ollama: {e}"
    except json.JSONDecodeError as e:
        return f"JSON Parsing Error: {e}, Raw Response: {response.text}"

def respond_to_chat(input: str, history: list[ChatMessage]):
    """Calls the Ollama API to get a response, including chat history."""
    return call_ollama_api(input, history)

def on_load(e: me.LoadEvent):
    me.set_theme_mode("system")

@me.page(
    security_policy=me.SecurityPolicy(
        allowed_iframe_parents=["https://google.github.io", "https://huggingface.co"]
    ),
    title="YUVI PlayGround for LLM-FineTuning",
    path="/",
    on_load=on_load
)
def page():
    state = me.state(State)

    with me.box(
        style=me.Style(
            background=me.theme_var("surface-container-lowest"),
            display="flex",
            flex_direction="column",
            height="100%",
        )
    ):
        with me.box(
            style=me.Style(
                display="flex", flex_direction="row", flex_grow=1, overflow="hidden"
            )
        ):
            with me.box(
                style=me.Style(
                    background=me.theme_var("surface-container-low"),
                    display="flex",
                    flex_direction="column",
                    flex_shrink=0,
                    position="absolute"
                    if state.sidebar_expanded and _is_mobile()
                    else None,
                    height="100%" if state.sidebar_expanded and _is_mobile() else None,
                    width=300 if state.sidebar_expanded else None,
                    z_index=2000,
                )
            ):
                sidebar()

            with me.box(
                style=me.Style(
                    display="flex",
                    flex_direction="column",
                    flex_grow=1,
                    padding=me.Padding(left=60)
                    if state.sidebar_expanded and _is_mobile()
                    else None,
                )
            ):
                header()
                with me.box(style=me.Style(flex_grow=1, overflow_y="scroll")):
                    if state.output:
                        chat_pane()
                    else:
                        examples_pane()
                chat_input()

def sidebar():
    state = me.state(State)
    with me.box(
        style=me.Style(
            display="flex",
            flex_direction="column",
            flex_grow=1,
        )
    ):
        with me.box(style=me.Style(display="flex", gap=20)):
            menu_icon(icon="menu", tooltip="Menu", on_click=on_click_menu_icon)
            if state.sidebar_expanded:
                me.text(
                    _APP_TITLE,
                    style=me.Style(margin=me.Margin(bottom=0, top=14)),
                    type="headline-6",
                )

        if state.sidebar_expanded:
            menu_item(icon="add", label="New chat", on_click=on_click_new_chat)
        else:
            menu_icon(icon="add", tooltip="New chat", on_click=on_click_new_chat)

        if state.sidebar_expanded:
            history_pane()

def history_pane():
    state = me.state(State)
    for index, chat in enumerate(state.history):
        with me.box(
            key=f"chat-{index}",
            on_click=on_click_history,
            style=me.Style(
                background=me.theme_var("surface-container"),
                border=me.Border.all(
                    me.BorderSide(
                        width=1, color=me.theme_var("outline-variant"), style="solid"
                    )
                ),
                border_radius=5,
                cursor="pointer",
                margin=me.Margin.symmetric(horizontal=10, vertical=10),
                padding=me.Padding.all(10),
                text_overflow="ellipsis",
            ),
        ):
            me.text(_truncate_text(chat[0]["content"]))

def header():
    state = me.state(State)
    with me.box(
        style=me.Style(
            align_items="center",
            background=me.theme_var("surface-container-lowest"),
            display="flex",
            gap=5,
            justify_content="space-between",
            padding=me.Padding.symmetric(horizontal=20, vertical=10),
            font_family="Roboto"
        )
    ):
        with me.box(style=me.Style(display="flex", gap=5)):
            if not state.sidebar_expanded:
                me.text(
                    _APP_TITLE,
                    style=me.Style(margin=me.Margin(bottom=0)),
                    type="headline-6",
                )

        with me.box(style=me.Style(display="flex", gap=5)):
            icon_button(
                key="",
                icon="dark_mode" if me.theme_brightness() == "light" else "light_mode",
                tooltip="Dark mode"
                if me.theme_brightness() == "light"
                else "Light mode",
                on_click=on_click_theme_brightness,
            )

def examples_pane():
    with me.box(
        style=me.Style(
            margin=me.Margin.symmetric(horizontal="auto"),
            padding=me.Padding.all(15),
            width=f"min({_CHAT_MAX_WIDTH}, 100%)",
          
        )
    ):
        with me.box(style=me.Style(margin=me.Margin(top=25), font_size=40,font_weight=70)):
            me.text(_EMPTY_CHAT_MESSAGE)

        with me.box(
            style=me.Style(
                display="flex",
                flex_direction="column" if _is_mobile() else "row",
                gap=20,
                margin=me.Margin(top=25),
            )
        ):
            for index, query in enumerate(_EXAMPLE_USER_QUERIES):
                with me.box(
                    key=f"query-{index}",
                    on_click=on_click_example_user_query,
                    style=me.Style(
                        background=me.theme_var("surface-container-highest"),
                        border_radius=30,
                        padding=me.Padding.all(20),
                        cursor="pointer",
                        font_family="Roboto"
                    ),
                ):
                    me.text(query)

def chat_pane():
    state = me.state(State)
    with me.box(
        style=me.Style(
            background=me.theme_var("surface-container-lowest"),
            color=me.theme_var("on-surface"),
            display="flex",
            flex_direction="column",
            margin=me.Margin.symmetric(horizontal="auto"),
            padding=me.Padding.all(15),
            width=f"min({_CHAT_MAX_WIDTH}, 100%)",
            font_family="Roboto" #YR
        )
    ):
        for index, msg in enumerate(state.output):
            if msg.role == "user":
                user_message(message=msg)
            else:
                bot_message(message_index=index, message=msg)

        if state.in_progress:
            with me.box(key="scroll-to", style=me.Style(height=250)):
                pass

def user_message(*, message: ChatMessage):
    with me.box(
        style=me.Style(
            display="flex",
            gap=15,
            justify_content="end",
            margin=me.Margin.all(20),
        )
    ):
        with me.box(
            style=me.Style(
                background=me.theme_var("surface-container-low"),
                border_radius=10,
                color=me.theme_var("on-surface-variant"),
                padding=me.Padding.symmetric(vertical=0, horizontal=10),
                width="66%",
            )
        ):
            me.markdown(message.content)

def bot_message(*, message_index: int, message: ChatMessage):
    with me.box(style=me.Style(display="flex", gap=20, margin=me.Margin.all(1))):
        text_avatar(
            background=me.theme_var("primary"),
            color=me.theme_var("on-primary"),
            label=_BOT_AVATAR_LETTER,
        
        )

        # Bot message response
        with me.box(style=me.Style(display="flex", flex_direction="column")):
            me.markdown(
                message.content,
                style=me.Style(color=me.theme_var("on-surface")),
            )

            # Actions panel
            with me.box():
                icon_button(
                    key=f"thumb_up-{message_index}",
                    icon="thumb_up",
                    is_selected=message.rating == 1,
                    tooltip="Good response",
                    on_click=on_click_thumb_up,
                )
                icon_button(
                    key=f"thumb_down-{message_index}",
                    icon="thumb_down",
                    is_selected=message.rating == -1,
                    tooltip="Bad response",
                    on_click=on_click_thumb_down,
                )
                icon_button(
                    key=f"restart-{message_index}",
                    icon="restart_alt",
                    tooltip="Regenerate answer",
                    on_click=on_click_regenerate,
                )



def chat_input():
    state = me.state(State)
    with me.box(
        style=me.Style(
            background=me.theme_var("surface-container")
            if _is_mobile()
            else me.theme_var("surface-container"),
            border_radius=16,
            display="flex",
            margin=me.Margin.symmetric(horizontal="auto", vertical=15),
            padding=me.Padding.all(8),
            width=f"min({_CHAT_MAX_WIDTH}, 100%)",
        )
    ):
        with me.box(
            style=me.Style(
                flex_grow=1,
                font_family="Roboto"
            )
        ):
            me.native_textarea(
                autosize=True,
                key="chat_input",
                min_rows=4,
                on_blur=on_chat_input,
                shortcuts={
                    me.Shortcut(key="Enter"): on_submit_chat_msg,  # Handle Enter key
                    me.Shortcut(shift=True, key="Enter"): lambda e: None,  # Allow Shift+Enter for newline
                },
                placeholder="Enter your prompt",
                style=me.Style(
                    background=me.theme_var("surface-container")
                    if _is_mobile()
                    else me.theme_var("surface-container"),
                    border=me.Border.all(
                        me.BorderSide(style="none"),
                    ),
                    color=me.theme_var("on-surface-variant"),
                    outline="none",
                    overflow_y="auto",
                    padding=me.Padding(top=16, left=16),
                    width="100%",
                ),
                value=state.input,
            )
        with me.content_button(
            disabled=state.in_progress,
            on_click=on_click_submit_chat_msg,
            type="icon",
        ):
            me.icon("send")

@me.component
def text_avatar(*, label: str, background: str, color: str):
    me.text(
        label,
        style=me.Style(
            background=background,
            border_radius="40%",
            color=color,
            font_size=20,
            height=40,
            line_height="1",
            margin=me.Margin(top=16),
            padding=me.Padding(top=10),
            text_align="center",
            width="45px",
        ),
    )

@me.component
def icon_button(
    *,
    icon: str,
    tooltip: str,
    key: str = "",
    is_selected: bool = False,
    on_click: Callable | None = None,
):
    selected_style = me.Style(
        background=me.theme_var("surface-container-low"),
        color=me.theme_var("on-surface-variant"),
    )
    with me.tooltip(message=tooltip):
        with me.content_button(
            type="icon",
            key=key,
            on_click=on_click,
            style=selected_style if is_selected else None,
        ):
            me.icon(icon)

@me.component
def menu_icon(
    *, icon: str, tooltip: str, key: str = "", on_click: Callable | None = None
):
    with me.tooltip(message=tooltip):
        with me.content_button(
            key=key,
            on_click=on_click,
            style=me.Style(margin=me.Margin.all(10)),
            type="icon",
        ):
            me.icon(icon)

@me.component
def menu_item(
    *, icon: str, label: str, key: str = "", on_click: Callable | None = None
):
    with me.box(on_click=on_click):
        with me.box(
            style=me.Style(
                background=me.theme_var("surface-container-high"),
                border_radius=20,
                cursor="pointer",
                display="inline-flex",
                gap=10,
                line_height=1,
                margin=me.Margin.all(10),
                padding=me.Padding(top=10, left=10, right=20, bottom=10),
            ),
        ):
            me.icon(icon)
            me.text(label, style=me.Style(height=24, line_height="24px"))

# Event Handlers

def on_click_example_user_query(e: me.ClickEvent):
    """Populates the user input with the example query"""
    state = me.state(State)
    _, example_index = e.key.split("-")
    state.input = _EXAMPLE_USER_QUERIES[int(example_index)]
    me.focus_component(key="chat_input")

def on_click_thumb_up(e: me.ClickEvent):
    """Gives the message a positive rating"""
    state = me.state(State)
    _, msg_index = e.key.split("-")
    msg_index = int(msg_index)
    state.output[msg_index].rating = 1

def on_click_thumb_down(e: me.ClickEvent):
    """Gives the message a negative rating"""
    state = me.state(State)
    _, msg_index = e.key.split("-")
    msg_index = int(msg_index)
    state.output[msg_index].rating = -1

def on_click_new_chat(e: me.ClickEvent):
    """Resets messages and saves the current chat to history."""
    state = me.state(State)
    if state.output:
        # Save the current chat to history
        state.history.insert(0, [asdict(message) for message in state.output])
    # Reset the output for a new chat
    state.output = []
    me.focus_component(key="chat_input")

def on_click_history(e: me.ClickEvent):
    """Loads existing chat from history and saves current chat."""
    state = me.state(State)
    _, chat_index = e.key.split("-")
    chat_index = int(chat_index)

    # Save the current chat to history if it exists
    if state.output:
        state.history.insert(0, [asdict(message) for message in state.output])

    # Load the selected chat from history
    chat_messages = [ChatMessage(**chat) for chat in state.history.pop(chat_index)]
    state.output = chat_messages
    me.focus_component(key="chat_input")

def on_click_theme_brightness(e: me.ClickEvent):
    """Toggles dark mode."""
    if me.theme_brightness() == "light":
        me.set_theme_mode("dark")
    else:
        me.set_theme_mode("light")

def on_click_menu_icon(e: me.ClickEvent):
    """Expands and collapses sidebar menu."""
    state = me.state(State)
    state.sidebar_expanded = not state.sidebar_expanded

def on_chat_input(e: me.InputBlurEvent):
    """Capture chat text input on blur."""
    state = me.state(State)
    state.input = e.value

def on_click_regenerate(e: me.ClickEvent):
    """Regenerates response from an existing message"""
    state = me.state(State)
    _, msg_index = e.key.split("-")
    msg_index = int(msg_index)

    # Get the user message which is the previous message
    user_message = state.output[msg_index - 1]
    # Get bot message to be regenerated
    assistant_message = state.output[msg_index]
    assistant_message.content = ""
    state.in_progress = True
    yield

    start_time = time.time()
    # Send in the old user input and chat history to get the bot response.
    # We make sure to only pass in the chat history up to this message.
    ollama_response = respond_to_chat(
        user_message.content, state.output[:msg_index]
    )
    assistant_message.content = ollama_response
    yield

    state.in_progress = False
    me.focus_component(key="chat_input")
    yield

def transform(input: str, history: list[mel.ChatMessage]):
    """
    Takes user input, sends it to the Ollama API, and yields the complete response.
    """
    # Call Ollama API with user input
    ollama_response = respond_to_chat(input, history)

    yield ollama_response

def on_submit_chat_msg(e: me.TextareaShortcutEvent):
    """Handles submitting a chat message when Enter is pressed."""
    state = me.state(State)
    state.input = e.value
    yield
    yield from _submit_chat_msg()
def on_click_submit_chat_msg(e: me.ClickEvent):
    yield from _submit_chat_msg()

def _submit_chat_msg():
    """Handles submitting a chat message."""
    state = me.state(State)
    if state.in_progress or not state.input:
        return
    input = state.input
    # Clear the text input.
    state.input = ""
    yield

    output = state.output
    if output is None:
        output = []
    output.append(ChatMessage(role="user", content=input))
    state.in_progress = True
    me.scroll_into_view(key="scroll-to")
    yield

    # Send user input and chat history to get the bot response.
    ollama_response = respond_to_chat(input, state.output)
    assistant_message = ChatMessage(role="bot", content=ollama_response)
    output.append(assistant_message)
    state.output = output

    state.in_progress = False
    me.focus_component(key="chat_input")
    yield

# Helpers

def _is_mobile():
    return me.viewport_size().width < _MOBILE_BREAKPOINT

def _truncate_text(text, char_limit=100):
    """Truncates text that is too long."""
    if len(text) <= char_limit:
        return text
    truncated_text = text[:char_limit].rsplit(" ", 1)[0]
    return truncated_text.rstrip(".,!?;:") + "..."