# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/4-ChatApp.ipynb.

# %% auto 0
__all__ = ['ChatAppConfig', 'ChatApp', 'create_default_chat_app', 'example_responder']

# %% ../../nbs/4-ChatApp.ipynb 1
from typing import Optional, Callable, List, Dict, Any
from dataclasses import dataclass, field
from fasthtml.common import *
from monsterui.all import Theme, Container, ContainerT, Card, CardT, TextPresets
import asyncio
import inspect

from .session import SessionManager, InMemorySessionManager, ChatSession, Message
from .service import ChatService, Responder, ErrorHandler
from .renderer import ChatRenderer
from .cards import ChatCard
from .design_system import get_color, get_spacing, get_typography

# %% ../../nbs/4-ChatApp.ipynb 3
@dataclass
class ChatAppConfig:
    """Configuration for ChatApp."""

    # App metadata
    app_title: str = "Chat Application"
    page_title: str = "Chat"

    # Initial messages
    initial_messages_factory: Optional[Callable[[], List[Message]]] = None

    # Styling (using design system defaults)
    bg_color: str = None
    header_style: str = None
    container_style: Optional[str] = None

    def __post_init__(self):
        """Initialize with design system defaults if not provided."""
        if self.bg_color is None:
            self.bg_color = get_color("dark_bg")
        if self.header_style is None:
            self.header_style = f"text-align: center; padding: {get_spacing('md')}; color: {get_color('light_text')};"

    # WebSocket settings
    ws_endpoint: str = "/ws"
    chat_endpoint: str = "/chat"

    # FastHTML extensions and headers
    markdown_enabled: bool = True
    syntax_highlighting: bool = True
    highlight_langs: List[str] = field(
        default_factory=lambda: ["python", "javascript", "html", "css"]
    )

    # Spinner configuration
    spinner_css: Optional[str] = None

    def get_default_initial_messages(self) -> List[Message]:
        """Get default initial messages."""
        return [
            Message(role="User", content="Hi"),
            Message(
                role="Assistant",
                content="Hello! I'm your helpful assistant. How can I assist you today?",
            ),
        ]

    def get_spinner_style(self) -> str:
        """Get spinner CSS styles (using design system)."""
        if self.spinner_css:
            return self.spinner_css

        return f"""
        .spinner {{
            display: inline-block;
            width: 20px;
            height: 20px;
            border: 3px solid {get_color("spinner_light")};
            border-top-color: {get_color("light_text")};
            border-radius: 50%;
            animation: spin 1s linear infinite;
        }}

        @keyframes spin {{
            to {{ transform: rotate(360deg); }}
        }}
        """

# %% ../../nbs/4-ChatApp.ipynb 4
class ChatApp:
    """Main chat application composing all components."""

    def __init__(
        self,
        session_manager: SessionManager,
        chat_service: ChatService,
        renderer: ChatRenderer,
        config: Optional[ChatAppConfig] = None,
    ):
        """
        Initialize ChatApp with dependency injection.

        Args:
            session_manager: Manages chat sessions
            chat_service: Handles message processing
            renderer: Renders UI components
            config: Application configuration
        """
        self.session_manager = session_manager
        self.chat_service = chat_service
        self.renderer = renderer
        self.config = config or ChatAppConfig()

        # Create FastHTML app
        self.app = self._create_fasthtml_app()
        self._register_routes()

    def _create_fasthtml_app(self) -> FastHTML:
        """Create and configure FastHTML application."""
        headers = list(Theme.blue.headers())

        headers.extend(
            [
                Link(rel="preconnect", href="https://fonts.googleapis.com"),
                Link(rel="preconnect", href="https://fonts.gstatic.com", crossorigin="anonymous"),
                Link(
                    rel="stylesheet",
                    href="https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@400;500;600;700&display=swap",
                ),
            ]
        )

        # Add markdown support
        if self.config.markdown_enabled:
            headers.append(MarkdownJS())

        # Add syntax highlighting
        if self.config.syntax_highlighting:
            headers.append(HighlightJS(langs=self.config.highlight_langs))

        # Add spinner styles
        headers.append(Style(self.config.get_spinner_style()))

        headers.append(Style(self.renderer.get_styles()))

        return FastHTML(exts="ws", hdrs=tuple(headers))

    def _get_initial_messages(self) -> List[Message]:
        """Get initial messages for new sessions."""
        if self.config.initial_messages_factory:
            return self.config.initial_messages_factory()
        return self.config.get_default_initial_messages()

    def _register_routes(self):
        """Register HTTP and WebSocket routes."""

        @self.app.route(self.config.chat_endpoint)
        def chat():
            """Main chat interface."""
            initial_messages = self._get_initial_messages()

            return (
                Title(self.config.page_title),
                Meta(name="viewport", content="width=device-width, initial-scale=1.0"),
                Body(
                    Container(
                        Div(
                            Div(
                                P("PYLOGUE", cls="tracking-widest text-xs uppercase text-emerald-200/70"),
                                H1(self.config.app_title, cls="text-4xl md:text-5xl font-semibold tracking-tight"),
                                P(
                                    "Streaming chat UI powered by MonsterUI + FastHTML.",
                                    cls=TextPresets.muted_sm,
                                ),
                                cls="space-y-3 mb-10",
                            ),
                            Card(
                                Div(
                                    self.renderer.render_messages(initial_messages),
                                    cls="chat-scroll",
                                ),
                                footer=self.renderer.render_form(),
                                header=Div(
                                    H2("Live Conversation", cls="text-2xl font-semibold"),
                                    P("Messages update in real time as the agent streams.", cls=TextPresets.muted_sm),
                                    cls="space-y-2",
                                ),
                                body_cls="space-y-4",
                                cls=("chat-shell", CardT.default),
                            ),
                            cls="pylogue-layout",
                        ),
                        cls=("pylogue-container", "mt-10", ContainerT.xl),
                    ),
                    cls="pylogue-app",
                ),
            )

        # Register WebSocket route with message handler
        @self.app.ws(
            self.config.ws_endpoint, conn=self._on_connect, disconn=self._on_disconnect
        )
        async def ws_handler(msg: str, send, ws):
            """WebSocket message handler."""
            await self._handle_websocket_message(msg, send, ws)

    def _on_connect(self, ws, send):
        """Handle WebSocket connection."""
        session_id = str(id(ws))
        initial_messages = self._get_initial_messages()
        self.session_manager.create_session(
            session_id=session_id, initial_messages=initial_messages
        )

    def _on_disconnect(self, ws):
        """Handle WebSocket disconnection."""
        session_id = str(id(ws))
        self.session_manager.delete_session(session_id)

    async def _handle_websocket_message(self, msg: str, send, ws):
        """
        Handle incoming WebSocket message with streaming support:
        1. Add user message
        2. Add empty assistant message for streaming
        3. Stream response tokens and update message progressively
        4. Clear input
        """
        session_id = str(id(ws))
        session = self.session_manager.get_session(session_id)

        if session is None:
            # Recreate session if somehow missing
            initial_messages = self._get_initial_messages()
            session = self.session_manager.create_session(
                session_id=session_id, initial_messages=initial_messages
            )

        # Step 1: Add user message and render
        session.add_message("User", msg)
        await send(self.renderer.render_messages(session.get_messages()))

        # Step 2: Add empty assistant message for streaming
        assistant_msg = session.add_message("Assistant", "", pending=False)

        # Step 3: Stream response and update progressively
        try:
            response_chunks = []
            chunk_count = 0
            async for chunk in self.chat_service.process_message_stream(msg, session):
                response_chunks.append(chunk)
                chunk_count += 1
                # Update the assistant message with accumulated response
                full_response = "".join(response_chunks)
                session.update_message(
                    assistant_msg.id, content=full_response, pending=False
                )
                # Send updated message list to UI
                # print(f"📤 Sending chunk #{chunk_count}: {repr(chunk)}")  # Debug
                await send(self.renderer.render_messages(session.get_messages()))

        except Exception as e:
            # Handle errors
            error_msg = f"Error: {str(e)}"
            session.update_message(assistant_msg.id, content=error_msg, pending=False)
            await send(self.renderer.render_messages(session.get_messages()))

        # Step 4: Clear input field
        await send(self.renderer.render_input())

    def run(
        self, host: str = "0.0.0.0", port: int = 5001, reload: bool = False, **kwargs
    ):
        """
        Run the chat application.

        Args:
            host: Host to bind to
            port: Port to bind to
            reload: Enable auto-reload on code changes
            **kwargs: Additional arguments for uvicorn.run()
        """
        import uvicorn

        print(f'Link: http://{"localhost" if host=="0.0.0.0" else host}:{port}')
        uvicorn.run(self.app, host=host, port=port, reload=reload, **kwargs)

    def get_app(self) -> FastHTML:
        """Get the underlying FastHTML app for advanced usage."""
        return self.app

# %% ../../nbs/4-ChatApp.ipynb 6
def create_default_chat_app(
    responder: Responder,
    config: Optional[ChatAppConfig] = None,
    card: Optional[ChatCard] = None,
    session_manager: Optional[SessionManager] = None,
    error_handler: Optional[ErrorHandler] = None,
    context_provider: Optional[Callable] = None,
) -> ChatApp:
    """
    Factory function to create a ChatApp with sensible defaults.

    Args:
        responder: The AI/LLM responder function
        config: Optional app configuration
        card: Optional custom ChatCard for styling
        session_manager: Optional custom session manager
        error_handler: Optional custom error handler
        context_provider: Optional context provider for responder

    Returns:
        Configured ChatApp instance
    """
    # Use defaults if not provided
    session_manager = session_manager or InMemorySessionManager()
    chat_service = ChatService(
        responder=responder,
        error_handler=error_handler,
        context_provider=context_provider,
    )
    renderer = ChatRenderer(card=card)
    config = config or ChatAppConfig()

    return ChatApp(
        session_manager=session_manager,
        chat_service=chat_service,
        renderer=renderer,
        config=config,
    )

# %% ../../nbs/4-ChatApp.ipynb 8
# Example responder
async def example_responder(message: str, context=None) -> str:
    """Example echo responder with latency."""
    await asyncio.sleep(0.5)  # Simulate processing time
    return f"Echo: {message}"
