import gradio as gr
from util import get_logger
import uuid
from pathlib import Path

# Import CRISPR-GPT modules
from crisprgpt import entry
from crisprgpt.logic import GradioMachineStateClass, concurrent_gradio_state_machine

logger = get_logger(__name__)
Path("log").mkdir(exist_ok=True)


# -------------------------
# Session Management
# -------------------------

def initialize_session():
    """Initialize a new CRISPR-GPT session."""
    session_id = str(uuid.uuid4().hex)
    full_task_list = [entry.EntryState]

    state = GradioMachineStateClass(full_task_list=full_task_list)
    concurrent_gradio_state_machine.reset(state)

    try:
        init_messages = concurrent_gradio_state_machine.loop(None, state)
        history = [(None, msg) for msg in init_messages] if init_messages else []
    except Exception as e:
        logger.error(f"Initialization error: {e}")
        history = [(None, "CRISPR-GPT ready! How can I assist you with molecular biology?")]

    return history, state, session_id


def save_chat(history, session_id):
    """Save chat history to file."""
    try:
        to_save = ""
        for user_msg, bot_msg in history:
            if user_msg:
                to_save += f"### User: \n{user_msg}\n\n"
            if bot_msg:
                to_save += f"### CRISPR-GPT: \n{bot_msg}\n\n"

        file_path = f"log/{session_id}.txt"
        with open(file_path, "w", encoding="utf-8") as f:
            f.write(to_save)
        return file_path
    except Exception as e:
        logger.error(f"Save error: {e}")
        return ""


# -------------------------
# Chat Logic
# -------------------------

def chat_respond(message, history, state, session_id):
    """Process user input and return bot response."""
    try:
        if not message or not str(message).strip():
            return history, state, session_id, ""

        bot_messages = concurrent_gradio_state_machine.loop(message, state)
        if not bot_messages:
            return history, state, session_id, ""

        for bot_msg in bot_messages:
            history.append((message, str(bot_msg)))
            message = None  # Show user message only once

        save_chat(history, session_id)
        return history, state, session_id, ""

    except Exception as e:
        logger.error(f"Chat error: {e}")
        history.append((message, f"Error: {str(e)}"))
        return history, state, session_id, ""


def reset_chat():
    """Reset chat session."""
    history, state, session_id = initialize_session()
    return history, state, session_id, ""


# -------------------------
# Gradio UI
# -------------------------

# Custom CSS
custom_css = """
* {
    font-family: Arial, sans-serif !important;
}
[data-testid="chatbot"] * {
    font-family: Arial, sans-serif !important;
    font-size: 14px !important;
}
pre, code {
    font-family: Courier, monospace !important;
}
"""

with gr.Blocks(title="CRISPR-GPT", theme=gr.themes.Soft(), css=custom_css) as demo:
    gr.Markdown("# 🧬 CRISPR-GPT")
    gr.Markdown("*AI Assistant for gene editing*")

    chatbot = gr.Chatbot(
        height=600,
        show_copy_button=True,
        type="tuples"
    )

    with gr.Row():
        msg = gr.Textbox(
            placeholder="Ask about gene-editing...",
            container=False,
            scale=4
        )
        send_btn = gr.Button("Send", variant="primary")

    with gr.Row():
        reset_btn = gr.Button("🔄 Reset Session")

    # Hidden states per session
    state = gr.State()
    session_id = gr.State()

    # Event handlers
    send_btn.click(
        chat_respond,
        [msg, chatbot, state, session_id],
        [chatbot, state, session_id, msg],
    )
    msg.submit(
        chat_respond,
        [msg, chatbot, state, session_id],
        [chatbot, state, session_id, msg],
    )

    reset_btn.click(reset_chat, outputs=[chatbot, state, session_id, msg])

    # Initialize when page loads
    demo.load(initialize_session, outputs=[chatbot, state, session_id])


if __name__ == "__main__":
    demo.queue()
    demo.launch(
        share=True,
        allowed_paths=["log/"],
        show_api=False,
        inbrowser=False,
    )