"""Message routes."""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

from fastapi import APIRouter, HTTPException, Query
from pydantic_ai import FunctionToolCallEvent
from pydantic_ai.messages import (
    PartDeltaEvent,
    PartStartEvent,
    TextPart as PydanticTextPart,
    TextPartDelta,
    ToolCallPart as PydanticToolCallPart,
)

from agentpool.agents.claude_code_agent.converters import derive_rich_tool_info
from agentpool.agents.events import (
    CompactionEvent,
    FileContentItem,
    LocationContentItem,
    StreamCompleteEvent,
    TextContentItem,
    ToolCallCompleteEvent,
    ToolCallProgressEvent,
    ToolCallStartEvent,
)
from agentpool.messaging.messages import ChatMessage
from agentpool.utils import identifiers as identifier
from agentpool.utils.pydantic_ai_helpers import safe_args_as_dict
from agentpool_server.opencode_server.converters import (
    _convert_params_for_ui,
    extract_user_prompt_from_parts,
    opencode_to_chat_message,
)
from agentpool_server.opencode_server.dependencies import StateDep  # noqa: TC001
from agentpool_server.opencode_server.models import (  # noqa: TC001
    AssistantMessage,
    MessagePath,
    MessageRequest,
    MessageTime,
    MessageUpdatedEvent,
    MessageWithParts,
    Part,
    PartUpdatedEvent,
    SessionCompactedEvent,
    SessionErrorEvent,
    SessionIdleEvent,
    SessionStatus,
    SessionStatusEvent,
    SessionUpdatedEvent,
    StepFinishPart,
    StepStartPart,
    TextPart,
    TimeCreated,
    TimeCreatedUpdated,
    TimeStartEnd,
    Tokens,
    TokensCache,
    ToolPart,
    ToolStateCompleted,
    ToolStateError,
    ToolStateRunning,
    UserMessage,
)
from agentpool_server.opencode_server.models.message import UserMessageModel
from agentpool_server.opencode_server.models.parts import (
    StepFinishTokens,
    TimeStart,
    TimeStartEndCompacted,
    TimeStartEndOptional,
    TokenCache,
)
from agentpool_server.opencode_server.routes.session_routes import (
    get_or_load_session,
    opencode_to_session_data,
)
from agentpool_server.opencode_server.time_utils import now_ms


if TYPE_CHECKING:
    from agentpool_server.opencode_server.state import ServerState


def _warmup_lsp_for_files(state: ServerState, file_paths: list[str]) -> None:
    """Warm up LSP servers for the given file paths.

    This starts LSP servers asynchronously based on file extensions.
    Like OpenCode's LSP.touchFile(), this triggers server startup without waiting.

    Args:
        state: Server state with LSP manager
        file_paths: List of file paths that were accessed
    """
    import logging

    logging.getLogger(__name__)
    print(f"[LSP] _warmup_lsp_for_files called with: {file_paths}")

    try:
        lsp_manager = state.get_or_create_lsp_manager()
        print("[LSP] Got LSP manager successfully")
    except RuntimeError as e:
        # No execution environment available for LSP
        print(f"[LSP] No LSP manager: {e}")
        return

    async def warmup_files() -> None:
        """Start LSP servers for each file path."""
        print("[LSP] warmup_files task started")
        from agentpool_server.opencode_server.models.events import LspUpdatedEvent

        servers_started = False
        for path in file_paths:
            # Find appropriate server for this file
            server_info = lsp_manager.get_server_for_file(path)
            print(f"[LSP] Server for {path}: {server_info.id if server_info else None}")
            if server_info is None:
                continue

            server_id = server_info.id
            if lsp_manager.is_running(server_id):
                print(f"[LSP] Server {server_id} already running")
                continue

            # Start server for workspace root
            root_uri = f"file://{state.working_dir}"
            try:
                print(f"[LSP] Starting server {server_id}...")
                await lsp_manager.start_server(server_id, root_uri)
                servers_started = True
                print(f"[LSP] Server {server_id} started successfully")
            except Exception as e:  # noqa: BLE001
                # Don't fail on LSP startup errors
                print(f"[LSP] Failed to start server {server_id}: {e}")

        # Emit lsp.updated event if any servers started
        if servers_started:
            print("[LSP] Broadcasting LspUpdatedEvent")
            await state.broadcast_event(LspUpdatedEvent.create())
        print("[LSP] warmup_files task completed")

    # Run warmup in background (don't block the event handler)
    print("[LSP] Creating background task for warmup")
    state.create_background_task(warmup_files(), name="lsp-warmup")


async def persist_message_to_storage(
    state: ServerState,
    msg: MessageWithParts,
    session_id: str,
) -> None:
    """Persist an OpenCode message to storage.

    Converts the OpenCode MessageWithParts to ChatMessage and saves it.

    Args:
        state: Server state with pool reference
        msg: OpenCode message to persist
        session_id: Session/conversation ID
    """
    if state.pool.storage is None:
        return

    try:
        # Convert to ChatMessage
        chat_msg = opencode_to_chat_message(msg, conversation_id=session_id)
        # Persist via storage manager
        await state.pool.storage.log_message(chat_msg)
    except Exception:  # noqa: BLE001
        # Don't fail the request if storage fails
        pass


async def _generate_session_title(
    state: ServerState,
    session_id: str,
    user_prompt: str,
    assistant_response: str,
) -> None:
    """Generate a title for the session in the background."""
    try:
        if not state.pool.storage:
            return

        # Create ChatMessage objects for the title generator
        messages = [
            ChatMessage[str](role="user", content=user_prompt),
            ChatMessage[str](role="assistant", content=assistant_response),
        ]

        # Generate title using storage manager
        title = await state.pool.storage.generate_conversation_title(
            messages=messages,
            conversation_id=session_id,
        )

        if title and session_id in state.sessions:
            # Update session with new title
            session = state.sessions[session_id]
            updated_session = session.model_copy(update={"title": title})
            state.sessions[session_id] = updated_session
            session_data = opencode_to_session_data(
                updated_session,
                agent_name=state.agent.name,
                pool_id=state.pool.manifest.config_file_path,
            )
            await state.pool.sessions.store.save(session_data)

            # Broadcast session update
            await state.broadcast_event(SessionUpdatedEvent.create(updated_session))
    except Exception:  # noqa: BLE001
        # Don't fail if title generation fails
        pass


router = APIRouter(prefix="/session/{session_id}", tags=["message"])


@router.get("/message")
async def list_messages(
    session_id: str,
    state: StateDep,
    limit: int | None = Query(default=None),
) -> list[MessageWithParts]:
    """List messages in a session."""
    session = await get_or_load_session(state, session_id)
    if session is None:
        raise HTTPException(status_code=404, detail="Session not found")

    messages = state.messages.get(session_id, [])
    if limit:
        messages = messages[-limit:]
    return messages


@router.post("/message")
async def send_message(  # noqa: PLR0915
    session_id: str,
    request: MessageRequest,
    state: StateDep,
) -> MessageWithParts:
    """Send a message and get response from the agent."""
    session = await get_or_load_session(state, session_id)
    if session is None:
        raise HTTPException(status_code=404, detail="Session not found")

    now = now_ms()
    # Create user message with sortable ID
    user_msg_id = identifier.ascending("message", request.message_id)
    user_message = UserMessage(
        id=user_msg_id,
        session_id=session_id,
        time=TimeCreated(created=now),
        agent=request.agent or "default",
        model=UserMessageModel(
            provider_id=request.model.provider_id if request.model else "agentpool",
            model_id=request.model.model_id if request.model else "default",
        ),
    )

    # Create parts from request
    user_parts: list[Part] = [
        TextPart(
            id=identifier.ascending("part"),
            message_id=user_msg_id,
            session_id=session_id,
            text=part.text,
        )
        for part in request.parts
        if part.type == "text"
    ]
    user_msg_with_parts = MessageWithParts(info=user_message, parts=user_parts)
    state.messages[session_id].append(user_msg_with_parts)
    # Persist user message to storage
    await persist_message_to_storage(state, user_msg_with_parts, session_id)
    # Broadcast user message created event
    await state.broadcast_event(MessageUpdatedEvent.create(user_message))
    # Broadcast user message parts so they appear in UI
    for part in user_parts:
        await state.broadcast_event(PartUpdatedEvent.create(part))
    state.session_status[session_id] = SessionStatus(type="busy")
    await state.broadcast_event(SessionStatusEvent.create(session_id, SessionStatus(type="busy")))
    # Extract user prompt text
    user_prompt = extract_user_prompt_from_parts([p.model_dump() for p in request.parts])
    # Create assistant message with sortable ID (must come after user message)
    assistant_msg_id = identifier.ascending("message")
    tokens = Tokens(cache=TokensCache(read=0, write=0))
    assistant_message = AssistantMessage(
        id=assistant_msg_id,
        session_id=session_id,
        parent_id=user_msg_id,  # Link to user message
        model_id=request.model.model_id if request.model else "default",
        provider_id=request.model.provider_id if request.model else "agentpool",
        mode=request.agent or "default",
        agent=request.agent or "default",
        path=MessagePath(cwd=state.working_dir, root=state.working_dir),
        time=MessageTime(created=now, completed=None),
        tokens=tokens,
        cost=0.0,
    )
    # Initialize assistant message with empty parts
    assistant_msg_with_parts = MessageWithParts(info=assistant_message, parts=[])
    state.messages[session_id].append(assistant_msg_with_parts)
    # Broadcast assistant message created
    await state.broadcast_event(MessageUpdatedEvent.create(assistant_message))
    # Add step-start part
    step_start = StepStartPart(
        id=identifier.ascending("part"),
        message_id=assistant_msg_id,
        session_id=session_id,
    )
    assistant_msg_with_parts.parts.append(step_start)
    await state.broadcast_event(PartUpdatedEvent.create(step_start))
    # Call the agent
    response_text = ""
    input_tokens = 0
    output_tokens = 0
    total_cost = 0.0  # Cost in dollars
    tool_parts: dict[str, ToolPart] = {}  # Track tool parts by call_id
    tool_outputs: dict[str, str] = {}  # Track accumulated output per tool call
    tool_inputs: dict[str, dict[str, Any]] = {}  # Track inputs per tool call
    # Track streaming text part for incremental updates
    text_part: TextPart | None = None
    text_part_id: str | None = None

    try:
        # Get the specified agent from the pool, or fall back to default
        agent = state.agent
        if request.agent and state.agent.agent_pool is not None:
            agent = state.agent.agent_pool.all_agents.get(request.agent, state.agent)

        # Stream events from the agent
        async for event in agent.run_stream(user_prompt):
            match event:
                # Text streaming start
                case PartStartEvent(part=PydanticTextPart(content=delta)):
                    response_text = delta
                    text_part_id = identifier.ascending("part")
                    text_part = TextPart(
                        id=text_part_id,
                        message_id=assistant_msg_id,
                        session_id=session_id,
                        text=delta,
                    )
                    assistant_msg_with_parts.parts.append(text_part)
                    await state.broadcast_event(PartUpdatedEvent.create(text_part, delta=delta))

                # Text streaming delta
                case PartDeltaEvent(delta=TextPartDelta(content_delta=delta)) if delta:
                    response_text += delta
                    if text_part is not None:
                        text_part = TextPart(
                            id=text_part.id,
                            message_id=assistant_msg_id,
                            session_id=session_id,
                            text=response_text,
                        )
                        # Update in parts list
                        for i, p in enumerate(assistant_msg_with_parts.parts):
                            if isinstance(p, TextPart) and p.id == text_part.id:
                                assistant_msg_with_parts.parts[i] = text_part
                                break
                        await state.broadcast_event(PartUpdatedEvent.create(text_part, delta=delta))

                # Tool call start - from Claude Code agent or toolsets
                case ToolCallStartEvent(
                    tool_name=tool_name,
                    tool_call_id=tool_call_id,
                    raw_input=raw_input,
                    title=title,
                ):
                    # Convert param names for OpenCode TUI compatibility
                    ui_input = _convert_params_for_ui(raw_input) if raw_input else {}
                    if tool_call_id in tool_parts:
                        # Update existing part with the custom title
                        existing = tool_parts[tool_call_id]
                        tool_inputs[tool_call_id] = ui_input or tool_inputs.get(tool_call_id, {})

                        updated = ToolPart(
                            id=existing.id,
                            message_id=existing.message_id,
                            session_id=existing.session_id,
                            tool=existing.tool,
                            call_id=existing.call_id,
                            state=ToolStateRunning(
                                status="running",
                                time=TimeStart(start=now_ms()),
                                input=tool_inputs[tool_call_id],
                                title=title,
                            ),
                        )
                        tool_parts[tool_call_id] = updated
                        for i, p in enumerate(assistant_msg_with_parts.parts):
                            if isinstance(p, ToolPart) and p.id == existing.id:
                                assistant_msg_with_parts.parts[i] = updated
                                break
                        await state.broadcast_event(PartUpdatedEvent.create(updated))
                    else:
                        # Create new tool part with the title
                        tool_inputs[tool_call_id] = ui_input
                        tool_outputs[tool_call_id] = ""
                        tool_state = ToolStateRunning(
                            status="running",
                            time=TimeStart(start=now_ms()),
                            input=ui_input,
                            title=title,
                        )
                        tool_part = ToolPart(
                            id=identifier.ascending("part"),
                            message_id=assistant_msg_id,
                            session_id=session_id,
                            tool=tool_name,
                            call_id=tool_call_id,
                            state=tool_state,
                        )
                        tool_parts[tool_call_id] = tool_part
                        assistant_msg_with_parts.parts.append(tool_part)
                        await state.broadcast_event(PartUpdatedEvent.create(tool_part))

                # Pydantic-ai tool call events (fallback for pydantic-ai agents)
                case (
                    FunctionToolCallEvent(part=tc_part)
                    | PartStartEvent(part=PydanticToolCallPart() as tc_part)
                ) if tc_part.tool_call_id not in tool_parts:
                    tool_call_id = tc_part.tool_call_id
                    tool_name = tc_part.tool_name
                    raw_input = safe_args_as_dict(tc_part)
                    # Convert param names for OpenCode TUI compatibility
                    ui_input = _convert_params_for_ui(raw_input)
                    # Store input and initialize output accumulator
                    tool_inputs[tool_call_id] = ui_input
                    tool_outputs[tool_call_id] = ""
                    # Derive initial title; toolset events may update it later
                    rich_info = derive_rich_tool_info(tool_name, raw_input)
                    tool_state = ToolStateRunning(
                        status="running",
                        time=TimeStart(start=now_ms()),
                        input=ui_input,
                        title=rich_info.title,
                    )
                    tool_part = ToolPart(
                        id=identifier.ascending("part"),
                        message_id=assistant_msg_id,
                        session_id=session_id,
                        tool=tool_name,
                        call_id=tool_call_id,
                        state=tool_state,
                    )
                    tool_parts[tool_call_id] = tool_part
                    assistant_msg_with_parts.parts.append(tool_part)
                    await state.broadcast_event(PartUpdatedEvent.create(tool_part))

                # Tool call progress
                case ToolCallProgressEvent(
                    tool_call_id=tool_call_id,
                    title=title,
                    items=items,
                    tool_name=tool_name,
                    tool_input=event_tool_input,
                ) if tool_call_id:
                    # Extract text content from items and accumulate
                    new_output = ""
                    file_paths: list[str] = []
                    for item in items:
                        if isinstance(item, TextContentItem):
                            new_output += item.text
                        elif isinstance(item, FileContentItem):
                            new_output += item.content
                            file_paths.append(item.path)
                        elif isinstance(item, LocationContentItem):
                            file_paths.append(item.path)

                    # Warm up LSP servers for accessed files (async, don't wait)
                    if file_paths:
                        _warmup_lsp_for_files(state, file_paths)

                    # Accumulate output (OpenCode streams via metadata.output)
                    if new_output:
                        tool_outputs[tool_call_id] = tool_outputs.get(tool_call_id, "") + new_output

                    if tool_call_id in tool_parts:
                        # Update existing part
                        existing = tool_parts[tool_call_id]
                        existing_title = getattr(existing.state, "title", "")
                        tool_input = tool_inputs.get(tool_call_id, {})
                        accumulated_output = tool_outputs.get(tool_call_id, "")
                        tool_state = ToolStateRunning(
                            status="running",
                            time=TimeStart(start=now_ms()),
                            title=title or existing_title,
                            input=tool_input,
                            metadata={"output": accumulated_output} if accumulated_output else None,
                        )
                        updated = ToolPart(
                            id=existing.id,
                            message_id=existing.message_id,
                            session_id=existing.session_id,
                            tool=existing.tool,
                            call_id=existing.call_id,
                            state=tool_state,
                        )
                        tool_parts[tool_call_id] = updated
                        for i, p in enumerate(assistant_msg_with_parts.parts):
                            if isinstance(p, ToolPart) and p.id == existing.id:
                                assistant_msg_with_parts.parts[i] = updated
                                break
                        await state.broadcast_event(PartUpdatedEvent.create(updated))
                    else:
                        # Create new tool part from progress event
                        ui_input = (
                            _convert_params_for_ui(event_tool_input) if event_tool_input else {}
                        )
                        tool_inputs[tool_call_id] = ui_input
                        accumulated_output = tool_outputs.get(tool_call_id, "")
                        tool_state = ToolStateRunning(
                            status="running",
                            time=TimeStart(start=now_ms()),
                            input=ui_input,
                            title=title or tool_name or "Running...",
                            metadata={"output": accumulated_output} if accumulated_output else None,
                        )
                        tool_part = ToolPart(
                            id=identifier.ascending("part"),
                            message_id=assistant_msg_id,
                            session_id=session_id,
                            tool=tool_name or "unknown",
                            call_id=tool_call_id,
                            state=tool_state,
                        )
                        tool_parts[tool_call_id] = tool_part
                        assistant_msg_with_parts.parts.append(tool_part)
                        await state.broadcast_event(PartUpdatedEvent.create(tool_part))

                # Tool call complete
                case ToolCallCompleteEvent(
                    tool_call_id=tool_call_id,
                    tool_result=result,
                ) if tool_call_id in tool_parts:
                    existing = tool_parts[tool_call_id]
                    result_str = str(result) if result else ""
                    tool_input = tool_inputs.get(tool_call_id, {})
                    is_error = isinstance(result, dict) and result.get("error")

                    if is_error:
                        new_state: ToolStateCompleted | ToolStateError = ToolStateError(
                            status="error",
                            error=str(result.get("error", "Unknown error")),
                            input=tool_input,
                            time=TimeStartEnd(start=now, end=now_ms()),
                        )
                    else:
                        new_state = ToolStateCompleted(
                            status="completed",
                            title=f"Completed {existing.tool}",
                            input=tool_input,
                            output=result_str,
                            time=TimeStartEndCompacted(start=now, end=now_ms()),
                        )

                    updated = ToolPart(
                        id=existing.id,
                        message_id=existing.message_id,
                        session_id=existing.session_id,
                        tool=existing.tool,
                        call_id=existing.call_id,
                        state=new_state,
                    )
                    tool_parts[tool_call_id] = updated
                    for i, p in enumerate(assistant_msg_with_parts.parts):
                        if isinstance(p, ToolPart) and p.id == existing.id:
                            assistant_msg_with_parts.parts[i] = updated
                            break
                    await state.broadcast_event(PartUpdatedEvent.create(updated))

                # Stream complete - extract token usage and cost
                case StreamCompleteEvent(message=msg) if msg:
                    if msg.usage:
                        input_tokens = msg.usage.input_tokens or 0
                        output_tokens = msg.usage.output_tokens or 0
                    if msg.cost_info and msg.cost_info.total_cost:
                        # Cost is in Decimal dollars, OpenCode expects float dollars
                        total_cost = float(msg.cost_info.total_cost)

                # Compaction event - emit session.compacted SSE event
                case CompactionEvent(session_id=compact_session_id, phase=phase):
                    if phase == "completed":
                        await state.broadcast_event(
                            SessionCompactedEvent.create(session_id=compact_session_id)
                        )

    except Exception as e:  # noqa: BLE001
        response_text = f"Error calling agent: {e}"
        # Emit session error event
        await state.broadcast_event(
            SessionErrorEvent.create(
                session_id=session_id,
                error_name=type(e).__name__,
                error_message=str(e),
            )
        )

    response_time = now_ms()

    # Create text part with response (only if we didn't stream it already)
    if response_text and text_part is None:
        text_part = TextPart(
            id=identifier.ascending("part"),
            message_id=assistant_msg_id,
            session_id=session_id,
            text=response_text,
            time=TimeStartEndOptional(start=now, end=response_time),
        )
        assistant_msg_with_parts.parts.append(text_part)

        # Broadcast text part update
        await state.broadcast_event(PartUpdatedEvent.create(text_part))
    elif text_part is not None:
        # Update the streamed text part with final timing
        final_text_part = TextPart(
            id=text_part.id,
            message_id=assistant_msg_id,
            session_id=session_id,
            text=response_text,
            time=TimeStartEndOptional(start=now, end=response_time),
        )
        # Update in parts list
        for i, p in enumerate(assistant_msg_with_parts.parts):
            if isinstance(p, TextPart) and p.id == text_part.id:
                assistant_msg_with_parts.parts[i] = final_text_part
                break

    step_finish = StepFinishPart(
        id=identifier.ascending("part"),
        message_id=assistant_msg_id,
        session_id=session_id,
        tokens=StepFinishTokens(
            cache=TokenCache(read=0, write=0),
            input=input_tokens,
            output=output_tokens,
            reasoning=0,
        ),
        cost=total_cost,
    )
    assistant_msg_with_parts.parts.append(step_finish)
    await state.broadcast_event(PartUpdatedEvent.create(step_finish))

    print(f"Response text: {response_text[:100] if response_text else 'EMPTY'}...")

    # Update assistant message with final timing and tokens
    updated_assistant = assistant_message.model_copy(
        update={
            "time": MessageTime(created=now, completed=response_time),
            "tokens": Tokens(
                cache=TokensCache(read=0, write=0),
                input=input_tokens,
                output=output_tokens,
                reasoning=0,
            ),
            "cost": total_cost,
        }
    )
    assistant_msg_with_parts.info = updated_assistant

    # Broadcast final message update
    await state.broadcast_event(MessageUpdatedEvent.create(updated_assistant))
    # Persist assistant message to storage
    await persist_message_to_storage(state, assistant_msg_with_parts, session_id)
    # Mark session as not running
    state.session_status[session_id] = SessionStatus(type="idle")
    await state.broadcast_event(SessionStatusEvent.create(session_id, SessionStatus(type="idle")))
    await state.broadcast_event(SessionIdleEvent.create(session_id))

    # Update session timestamp
    session = state.sessions[session_id]
    state.sessions[session_id] = session.model_copy(
        update={"time": TimeCreatedUpdated(created=session.time.created, updated=response_time)}
    )
    # Trigger title generation if session has default title
    if session.title == "New Session" and state.pool.storage:
        # Convert user_prompt to string if it's a list
        prompt_str = user_prompt if isinstance(user_prompt, str) else str(user_prompt)
        state.create_background_task(
            _generate_session_title(state, session_id, prompt_str, response_text),
            name=f"generate_title_{session_id}",
        )
    return assistant_msg_with_parts


@router.get("/message/{message_id}")
async def get_message(
    session_id: str,
    message_id: str,
    state: StateDep,
) -> MessageWithParts:
    """Get a specific message."""
    session = await get_or_load_session(state, session_id)
    if session is None:
        raise HTTPException(status_code=404, detail="Session not found")

    for msg in state.messages.get(session_id, []):
        if msg.info.id == message_id:
            return msg

    raise HTTPException(status_code=404, detail="Message not found")
