import streamlit as st
from chat_langraph import system, workflow, HumanMessage, AIMessage, get_all_chat_ids, ToolMessage
import uuid
import os
import base64

st.set_page_config(layout="wide")
st.title("College Chatbot")

TEMP_DIR = "/tmp"
os.makedirs(TEMP_DIR, exist_ok=True)


# --- Helpers ---

def set_title(messages):
    if messages:
        title = "New Chat"
        st.session_state.chat_dict[st.session_state.current_chat_id] = title


def set_config():
    return {"configurable": {"thread_id": st.session_state.current_chat_id}}


def load_session_state():
    if "chats" not in st.session_state:
        st.session_state.chats = get_all_chat_ids()
    if "current_chat_id" not in st.session_state:
        if len(st.session_state.chats) > 0:
            st.session_state.current_chat_id = st.session_state.chats[-1]
        else:
            new_id = str(uuid.uuid4())
            st.session_state.chats.append(new_id)
            st.session_state.current_chat_id = new_id
    if "chat_dict" not in st.session_state:
        st.session_state.chat_dict = {}


def render_sidebar():
    with st.sidebar:
        st.title("Chats")
        if st.button("➕ New Chat"):
            new_id = str(uuid.uuid4())
            st.session_state.chats.append(new_id)
            st.session_state.current_chat_id = new_id
            config = {"configurable": {"thread_id": new_id}}
            workflow.update_state(config, {"messages": [system]})
            st.session_state.chat_dict[new_id] = "New Chat"

        for chat_id in st.session_state.chats:
            if st.button(st.session_state.chat_dict.get(chat_id, "New Chat"), key=chat_id):
                st.session_state.current_chat_id = chat_id


def create_download_link(file_path: str, label: str = None) -> str:
    """Generate HTML download link for a file."""
    if not os.path.exists(file_path):
        return ""
    try:
        with open(file_path, "rb") as f:
            data = f.read()
        b64 = base64.b64encode(data).decode()
        label = label or f"📥 Download {os.path.basename(file_path)}"
        href = f'<a href="data:file/octet-stream;base64,{b64}" download="{os.path.basename(file_path)}">{label}</a>'
        return href
    except Exception as e:
        return f"Error creating download link: {e}"


def show_file(file_path: str):
    """Show file inline + download link."""
    if not os.path.exists(file_path):
        return
    ext = os.path.splitext(file_path)[1].lower()
    if ext in [".png", ".jpg", ".jpeg"]:
        st.image(file_path, caption=os.path.basename(file_path))
    elif ext in [".txt", ".py", ".java", ".cpp", ".md"]:
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            st.code(f.read(), language=ext.lstrip("."))
    st.markdown(create_download_link(file_path), unsafe_allow_html=True)


def render_tool_message(tool_message: ToolMessage):
    """Render tool execution based on tool name instead of message content."""
    file_related_keywords = ["read", "write", "file", "save", "export", "create"]
    with st.chat_message("assistant"):
        tool_name = getattr(tool_message, "name", "").lower()
        st.info(f"🧰 Tool used: {tool_name or 'Unknown Tool'}")

        # Check if this is a file-related tool
        if any(k in tool_name for k in file_related_keywords):
            # Find all files in TEMP_DIR (freshly modified ones)
            created_files = sorted(
                [os.path.join(TEMP_DIR, f) for f in os.listdir(TEMP_DIR)],
                key=lambda x: os.path.getmtime(x),
                reverse=True,
            )
            if created_files:
                st.success("📄 File(s) created by tool:")
                for file_path in created_files[:3]:  # show up to 3 recent
                    show_file(file_path)
            else:
                st.warning("No new file detected in /tmp.")


def loadchats():
    if "current_chat_id" not in st.session_state:
        return []
    config = {"configurable": {"thread_id": st.session_state.current_chat_id}}
    state = workflow.get_state(config)
    messages = state.values.get("messages", [])
    for message in messages:
        if isinstance(message, HumanMessage):
            with st.chat_message("human"):
                st.write(message.content)
        elif isinstance(message, AIMessage):
            with st.chat_message("assistant"):
                st.write(message.content)
        elif isinstance(message, ToolMessage):
            render_tool_message(message)
    return messages


# --- Main Chat Flow ---

load_session_state()
render_sidebar()

if "current_chat_id" in st.session_state:
    loadchats()
    if user_input := st.chat_input("Your message:"):
        with st.chat_message("human"):
            st.write(user_input)

        with st.chat_message("assistant"):
            response_placeholder = st.empty()
            full_response = ""

            for message, metadata in workflow.stream(
                {"messages": [system, HumanMessage(user_input)]},
                config={"configurable": {"thread_id": st.session_state.current_chat_id}},
                stream_mode="messages",
            ):
                if isinstance(message, AIMessage):
                    full_response += message.content or ""
                elif isinstance(message, ToolMessage):
                    render_tool_message(message)

            response_placeholder.markdown(full_response)