import streamlit as st
from datetime import datetime
import os
import json
from pathlib import Path
from json import JSONDecodeError
from FreeLLM import HuggingChatAPI, ChatGPTAPI, BingChatAPI

# Page Configuration
st.set_page_config(
    page_title="FREE AUTOHUGGS 🤗",
    page_icon="🤗",
    layout="wide",
    menu_items={
        "Get help": "getkyklos@dontsp.am",
        "About": "# *🤗 FREE AUTOHUGGS 🤗*"
    }
)

# Modern Dark Theme CSS
st.markdown("""
<style>
    /* Base theme */
    .main {
        background-color: #0e1117;
        color: #e0e0e0;
    }
    
    /* Sidebar styling */
    .css-1d391kg {
        background-color: #1e1e2e;
    }
    
    /* Input fields */
    .stTextInput input, .stTextArea textarea, .stNumberInput input {
        background-color: #262630;
        color: #ffffff;
        border: 1px solid #464660;
        border-radius: 8px;
        padding: 12px;
    }
    
    /* Buttons */
    .stButton>button {
        background: linear-gradient(90deg, #ff4b4b 0%, #ff6b6b 100%);
        color: white;
        border: none;
        padding: 0.6em 1.2em;
        border-radius: 8px;
        font-weight: 600;
        transition: all 0.3s ease;
    }
    
    .stButton>button:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 12px rgba(255, 75, 75, 0.3);
    }
    
    /* Chat messages */
    .chat-message {
        padding: 1rem;
        border-radius: 12px;
        margin: 1rem 0;
        box-shadow: 0 2px 8px rgba(0,0,0,0.1);
    }
    
    .user-message {
        background-color: #2a2a3a;
        border-left: 4px solid #ff4b4b;
    }
    
    .assistant-message {
        background-color: #323242;
        border-left: 4px solid #4b8bff;
    }
    
    /* Headers */
    h1, h2, h3 {
        color: #ffffff;
    }
</style>
""", unsafe_allow_html=True)

# Sidebar Configuration
with st.sidebar:
    st.header("🛠️ Agent Configuration")
    
    assistant_role_name = st.text_input("👤 Assistant Role", "Python Programmer")
    user_role_name = st.text_input("👥 User Role", "Stock Trader")
    task = st.text_area("📋 Task", "Develop a trading bot for the stock market")
    word_limit = st.slider("🔤 Word Limit", 10, 1500, 50)
    
    st.markdown("---")
    
    st.header("⚙️ Advanced Options")
    llm_choice = st.selectbox("🤖 Select LLM API", ["HuggingChat", "ChatGPT", "BingChat"])
    auto_save = st.checkbox("💾 Auto-save conversation", value=True)
    transcript_filename = st.text_input("📄 Transcript filename", "conversation_transcript.txt")

# Cookie File Validation
COOKIE_FILE = "cookiesHuggingChat.json"
try:
    if not Path(COOKIE_FILE).exists():
        st.error(f"❌ File '{COOKIE_FILE}' not found! Please create it with your cookies in JSON format.")
        st.stop()
        
    with open(COOKIE_FILE, "r") as file:
        cookies = json.loads(file.read())
except JSONDecodeError:
    st.error(f"❌ Invalid JSON in '{COOKIE_FILE}'. Please check your file format!")
    st.stop()
except Exception as e:
    st.error(f"❌ Error reading cookie file: {str(e)}")
    st.stop()

# LLM API Initialization
def get_llm(api_choice: str):
    llm_map = {
        "HuggingChat": HuggingChatAPI.HuggingChat,
        "ChatGPT": ChatGPTAPI.ChatGPT,
        "BingChat": BingChatAPI.BingChat
    }
    
    if api_choice not in llm_map:
        st.error("❌ Invalid LLM API selection.")
        st.stop()
        
    try:
        return llm_map[api_choice](cookiepath=str(Path(COOKIE_FILE)))
    except Exception as e:
        st.error(f"❌ Error initializing {api_choice}: {str(e)}")
        st.stop()

llm = get_llm(llm_choice)

# Initialize session state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "conversation_start" not in st.session_state:
    st.session_state.conversation_start = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# Main chat interface
st.title("💬 AI Chat Interface")

# Display chat history
for message in st.session_state.chat_history:
    if message["role"] == "assistant":
        st.markdown(
            f"<div class='chat-message assistant-message'><strong>{assistant_role_name}:</strong> {message['content']}</div>",
            unsafe_allow_html=True
        )
    elif message["role"] == "user":
        st.markdown(
            f"<div class='chat-message user-message'><strong>{user_role_name}:</strong> {message['content']}</div>",
            unsafe_allow_html=True
        )
    else:
        st.markdown(
            f"<div class='chat-message'><em>{message['content']}</em></div>",
            unsafe_allow_html=True
        )

# Input area
user_input = st.text_input("💭 Your message:", key="user_input")

col1, col2, col3 = st.columns(3)

with col1:
    if st.button("📤 Send"):
        if user_input:
            # Add user message to history
            st.session_state.chat_history.append({"role": "user", "content": user_input})
            
            # Get AI response
            try:
                with st.spinner("🤔 AI is thinking..."):
                    response = llm(user_input)
                st.session_state.chat_history.append({"role": "assistant", "content": response})
                
                # Auto-save if enabled
                if auto_save:
                    with open(transcript_filename, "a", encoding="utf-8") as f:
                        f.write(f":User  {user_input}\n")
                        f.write(f"Assistant: {response}\n")
                        
                st.experimental_rerun()
            except Exception as e:
                st.error(f"❌ Error getting AI response: {str(e)}")
with col2:
    if st.button("📝 Clear Chat"):
        st.session_state.chat_history = []
        st.experimental_rerun()
with col3:
    if st.button("💾 Save Transcript"):
        try:
            transcript = "\n".join([
                f"{msg['role'].capitalize()}: {msg['content']}"
                for msg in st.session_state.chat_history
            ])
            with open(transcript_filename, "w", encoding="utf-8") as f:
                f.write(transcript)
            st.success("✅ Transcript saved successfully!")
        except Exception as e:
            st.error(f"❌ Error saving transcript: {str(e)}")
# Additional Features
# Function to display the current time
def display_time():
    current_time = datetime.now().strftime("%H:%M:%S")
    st.sidebar.markdown(f"🕒 Current Time: {current_time}")

# Call the function to display time
display_time()

# Function to handle user input and AI response
def handle_input(user_input):
    if user_input:
        st.session_state.chat_history.append({"role": "user", "content": user_input})
        try:
            with st.spinner("🤔 AI is thinking..."):
                response = llm(user_input)
            st.session_state.chat_history.append({"role": "assistant", "content": response})
            if auto_save:
                with open(transcript_filename, "a", encoding="utf-8") as f:
                    f.write(f":User  {user_input}\n")
                    f.write(f"Assistant: {response}\n")
            st.experimental_rerun()
        except Exception as e:
            st.error(f"❌ Error getting AI response: {str(e)}")

# Input area with button to send message
if st.button("📤 Send"):
    handle_input(user_input)

# Function to clear chat history
def clear_chat():
    st.session_state.chat_history = []
    st.experimental_rerun()

# Clear chat button
if st.button("📝 Clear Chat"):
    clear_chat()

# Function to save transcript
def save_transcript():
    try:
        transcript = "\n".join([
            f"{msg['role'].capitalize()}: {msg['content']}"
            for msg in st.session_state.chat_history
        ])
        with open(transcript_filename, "w", encoding="utf-8") as f:
            f.write(transcript)
        st.success("✅ Transcript saved successfully!")
    except Exception as e:
        st.error(f"❌ Error saving transcript: {str(e)}")

# Save transcript button
if st.button("💾 Save Transcript"):
    save_transcript()