import logging
from typing import Final, List, Union

import gradio as gr
from gradio import ChatMessage

from agent import BasicAgent

MAX_INPUT_CHARS: Final[int] = 4000

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

agent = BasicAgent()


def submit(user_message: str, history: List[Union[ChatMessage, dict]]):
    history = history or []

    logger.info(f"User Message: {user_message}")

    if not user_message:
        return history + [
            ChatMessage(role="user", content=user_message),
            ChatMessage(role="assistant", content="Input too short."),
        ], ""

    if len(user_message) > MAX_INPUT_CHARS:
        return history + [
            ChatMessage(role="user", content=user_message),
            ChatMessage(
                role="assistant",
                content=f"Input too long; please shorten to under {MAX_INPUT_CHARS} characters.",
            ),
        ], ""

    # Add user message to history
    history = history + [ChatMessage(role="user", content=user_message)]

    # Build prompt from last messages
    last_n = 6
    prompt_parts = []
    for msg in history[-last_n:]:
        # Handle both ChatMessage objects and dicts
        if isinstance(msg, dict):
            role = msg.get("role", "user")
            content = msg.get("content", "")
        else:
            role = msg.role
            content = msg.content

        prefix = "User:" if role == "user" else "Assistant:"
        prompt_parts.append(f"{prefix} {content}")

    prompt = "\n".join(prompt_parts) + "\nAssistant:"

    try:
        answer = agent(prompt)
        logger.info(f"Agent Response: {answer}")
    except Exception as e:
        logger.error(f"Agent error: {e}")
        answer = "Sorry, I encountered an error. Please try again."

    # Add agent response to history
    history = history + [ChatMessage(role="assistant", content=str(answer))]
    logger.info(f"Updated history: {history}")

    return history, ""


with gr.Blocks(title="BasicAgent Chat") as demo:
    chatbot = gr.Chatbot()
    message = gr.Textbox()
    clear_btn = gr.ClearButton([message, chatbot])
    submit_btn = gr.Button("Send")

    message.submit(
        submit,
        inputs=[message, chatbot],
        outputs=[chatbot, message],
    )

    submit_btn.click(submit, inputs=[message, chatbot], outputs=[chatbot, message])
    clear_btn.click(lambda: [], None, chatbot)


if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=7860)