import uuid
from datetime import datetime
from typing import Callable, Dict, Iterator, List, Optional

import streamlit as st


def handle_user_input(
        generate_response: Optional[Callable],
        generate_response_stream: Optional[Callable[[List[Dict]], Iterator[str]]],
        generate_title: Optional[Callable],
        count_tokens: Optional[Callable],
        chat_placeholder: str,
        spinner_text: str,
        max_history_tokens: Optional[int],
        save_conversations_func: Callable,
        user_avatar: Optional[str] = None,
        assistant_avatar: Optional[str] = None,
        chat_input_value: Optional[str] = None,
):
    """
    Process user input, generate responses, and update conversation.

    Args:
        generate_response: Function to generate responses to user messages
        generate_response_stream: Optional stream generator (takes priority over generate_response)
        generate_title: Function to generate conversation titles
        count_tokens: Function to count tokens in conversations
        chat_placeholder: Placeholder text for the chat input
        spinner_text: Text to display while generating a response
        max_history_tokens: Maximum number of tokens to keep in history
        save_conversations_func: Function to save conversations
        user_avatar: Optional avatar for user messages
        assistant_avatar: Optional avatar for assistant messages
    """
    from .ui_components import DEFAULT_USER_AVATAR, DEFAULT_ASSISTANT_AVATAR
    user_avatar = user_avatar or DEFAULT_USER_AVATAR
    assistant_avatar = assistant_avatar or DEFAULT_ASSISTANT_AVATAR
    # Use chat_input_value from chat_ui (single widget, avoids duplicate input box)
    prompt = st.session_state.pop("pending_suggestion", None) or chat_input_value
    if prompt:
        # Add user message to chat (continue in current conversation; create_conversation_if_needed handles new convo when needed)
        with st.chat_message("user", avatar=user_avatar):
            st.markdown(prompt)

        # Add user message to state
        st.session_state.messages.append({"role": "user", "content": prompt})

        # Create a conversation if needed
        create_conversation_if_needed(generate_title, prompt, save_conversations_func)

        # Flag to track if title was updated
        title_updated = False

        # Update title if this is the first message and title is still default
        convo_id = st.session_state.current_convo_id
        if (len(st.session_state.messages) == 1 and
                st.session_state.conversations[convo_id]["title"] == "New Conversation" and
                generate_title):
            st.session_state.conversations[convo_id]["title"] = generate_title(prompt)
            title_updated = True

        # Prefer streaming over non-streaming when both could be provided
        responder = generate_response_stream or generate_response
        if responder:
            # Truncate messages if needed before calling
            if max_history_tokens and count_tokens:
                st.session_state.messages = truncate_messages(
                    st.session_state.messages,
                    max_history_tokens,
                    count_tokens
                )

            if generate_response_stream and responder == generate_response_stream:
                with st.chat_message("assistant", avatar=assistant_avatar):
                    response = st.write_stream(generate_response_stream(st.session_state.messages))
            else:
                with st.chat_message("assistant", avatar=assistant_avatar):
                    with st.spinner(spinner_text):
                        response = generate_response(st.session_state.messages)
                    st.markdown(response)

            # Add response to state and save
            st.session_state.messages.append({"role": "assistant", "content": response})

            # Update conversation in storage
            st.session_state.conversations[convo_id]["messages"] = st.session_state.messages

            # Update token count if provided
            if count_tokens:
                token_count = count_tokens(st.session_state.messages)
                st.session_state.conversations[convo_id]["token_count"] = token_count

            save_conversations_func()

            # Rerun so welcome screen is replaced by chat history (and title update is visible)
            st.rerun()


def create_conversation_if_needed(
        generate_title: Optional[Callable], 
        first_prompt: str,
        save_conversations_func: Callable
):
    """
    Create a new conversation if one doesn't exist.
    
    Args:
        generate_title: Function to generate conversation titles
        first_prompt: The first user message
        save_conversations_func: Function to save conversations
    """
    if st.session_state.current_convo_id is None:
        convo_id = str(uuid.uuid4())
        st.session_state.current_convo_id = convo_id

        # Set default title, which may be updated later
        title = "New Conversation"
        if generate_title:
            title = generate_title(first_prompt)

        st.session_state.conversations[convo_id] = {
            "id": convo_id,
            "title": title,
            "messages": st.session_state.messages,
            "created_at": datetime.now().isoformat(),
            "token_count": 0
        }
        save_conversations_func()


def truncate_messages(messages: List[Dict], max_tokens: int, count_tokens: Callable) -> List[Dict]:
    """
    Truncate message history to fit within token limit.
    
    Args:
        messages: List of message dictionaries
        max_tokens: Maximum number of tokens to keep
        count_tokens: Function to count tokens
        
    Returns:
        Truncated list of messages
    """
    if not messages:
        return messages

    # Always keep the system message if present
    system_message = None
    other_messages = messages.copy()

    if messages[0]["role"] == "system":
        system_message = messages[0]
        other_messages = messages[1:]

    # Check if we need to truncate
    if count_tokens(messages) <= max_tokens:
        return messages

    # Truncate older messages first: keep system + most recent messages within limit
    truncated_messages = []
    for message in reversed(other_messages):
        candidate = ([system_message] if system_message else []) + [message] + truncated_messages
        if count_tokens(candidate) <= max_tokens:
            truncated_messages = [message] + truncated_messages
        else:
            break
    if system_message:
        truncated_messages = [system_message] + truncated_messages
    return truncated_messages


def create_system_message(system_prompt: str):
    """
    Create a system message and add it to the current conversation.
    
    Args:
        system_prompt: Content for the system message
    """
    if not system_prompt:
        return
    
    # Check if there's already a system message
    if st.session_state.messages and st.session_state.messages[0]["role"] == "system":
        # Update existing system message
        st.session_state.messages[0]["content"] = system_prompt
    else:
        # Add system message at the beginning
        st.session_state.messages.insert(0, {
            "role": "system",
            "content": system_prompt
        })
    
    # Update the conversation in storage if a conversation exists
    if st.session_state.current_convo_id:
        convo_id = st.session_state.current_convo_id
        st.session_state.conversations[convo_id]["messages"] = st.session_state.messages


def format_markdown_message(message_obj: Dict) -> str:
    """
    Format a message object as markdown.
    
    Args:
        message_obj: Message dictionary with role and content
        
    Returns:
        Formatted markdown string
    """
    role = message_obj.get("role", "unknown").upper()
    content = message_obj.get("content", "")
    
    return f"**{role}**:\n\n{content}\n\n"


def extract_code_blocks(message: str) -> List[Dict[str, str]]:
    """
    Extract code blocks from a message.
    
    Args:
        message: Message content potentially containing code blocks
        
    Returns:
        List of dictionaries with code and language
    """
    import re
    
    # Match code blocks with language specification
    # Format: ```language\ncode\n```
    pattern = r"```([a-zA-Z0-9_]*)\n([\s\S]*?)\n```"
    
    matches = re.findall(pattern, message)
    
    code_blocks = []
    for lang, code in matches:
        code_blocks.append({
            "language": lang.strip() if lang.strip() else "text",
            "code": code.strip()
        })
    
    return code_blocks
