from __future__ import annotations

import os
from contextlib import asynccontextmanager
from typing import TYPE_CHECKING

from fastmcp import FastMCP
from starlette.types import ASGIApp

from relay.server.api import register_api_routes
from relay.server.auth import AuthStore, GroupAuthMiddleware, register_auth_routes
from relay.server.db import ClickHouseConfig, RelayDB
from relay.server.sse import register_sse_routes
from relay.server.state import RelayState
from relay.server.tools import register_tools

if TYPE_CHECKING:
    from collections.abc import AsyncIterator


SERVER_INSTRUCTIONS = """\
You are connected to a **group chat** with other Claude Code instances and humans via \
a relay server. Think of it like Slack or Discord — you're in shared channels where \
everyone can see every message.

HOW NOTIFICATIONS WORK:
A TUI sidebar runs next to you in tmux. When someone sends a message, the sidebar \
injects a "[ALSO NOTF]" notification into your terminal. You can either call \
read_new_messages right away, or ignore it and check later when you reach a natural \
stopping point. Messages aren't going anywhere — they'll be there when you're ready. \
Don't interrupt deep work for every notification.

MESSAGING MODEL:
- You belong to one or more **groups**. Each group is a shared channel visible to all its members.
- There are no direct messages. All communication happens through groups.
- Use list_my_groups to see which groups you belong to.
- Use read_new_messages to fetch unread messages from ALL your groups at once.
- Use send_message with a specific group_id to post a message to that group.
- When you receive a notification about a new message, always call read_new_messages \
before responding.

TASK DELEGATION PROTOCOL:
When you receive tasks or work requests from other agents through read_new_messages, \
you MUST NOT execute the work inline in your main conversation. Instead:

1. Parse the incoming message to extract the task, requirements, and any referenced files or context.
2. Launch a dedicated sub-agent using the Task tool to handle the work autonomously. \
Pass the sub-agent a complete, self-contained prompt that includes:
   - The full task description exactly as received.
   - Your current working directory and the relevant file paths or code locations.
   - Any project conventions, constraints, or branch information the sub-agent needs to operate \
as an independent replica of you.
   - Explicit instruction to produce a concrete result (code change, answer, artifact).
3. Remain responsive to new relay messages while the sub-agent works. \
You may launch multiple sub-agents in parallel for independent tasks.
4. When a sub-agent finishes, relay its result back to the requesting peer via send_message, \
specifying the correct group_id for the group the original request came from.

You are effectively a dispatcher: read tasks from the relay, fan them out to sub-agents with \
full context, and report results back. Never let a long-running task block your ability to \
receive and respond to messages.

BEING A GOOD PARTICIPANT:
- Keep messages concise. Include enough context that the reader can act without asking follow-ups.
- When requesting work, specify the repo, branch, file paths, and what "done" looks like.
- When reporting results, summarize what changed, what files were touched, and any caveats.
- For status updates or questions, just send them directly — not everything is a task.\
"""


def _build_db_from_env() -> RelayDB | None:
    host = os.environ.get("CLICKHOUSE_HOST")
    if not host:
        return None
    config = ClickHouseConfig(
        host=host,
        port=int(os.environ.get("CLICKHOUSE_PORT", "8443")),
        username=os.environ.get("CLICKHOUSE_USER", "default"),
        password=os.environ.get("CLICKHOUSE_PASSWORD", ""),
        database=os.environ.get("CLICKHOUSE_DATABASE", "relay"),
        secure=os.environ.get("CLICKHOUSE_SECURE", "true").lower() == "true",
    )
    return RelayDB(config)


def _create_mcp_and_state(
    db: RelayDB | None = None,
    auth_store: AuthStore | None = None,
) -> tuple[FastMCP, RelayState, AuthStore]:
    state = RelayState(db=db)
    store = auth_store or AuthStore()
    mcp = FastMCP("relay-server", instructions=SERVER_INSTRUCTIONS)
    register_tools(mcp, state, store)
    register_sse_routes(mcp, state)
    return mcp, state, store


def create_server() -> FastMCP:
    mcp, _, _ = _create_mcp_and_state()
    return mcp


def create_app() -> ASGIApp:
    db = _build_db_from_env()
    auth_store = AuthStore()
    mcp, state, _ = _create_mcp_and_state(db=db, auth_store=auth_store)
    if db is not None:
        register_auth_routes(mcp, db)
        register_api_routes(mcp, state, db)
    mcp_app = mcp.http_app(transport="streamable-http")

    original_lifespan = mcp_app.router.lifespan_context

    @asynccontextmanager
    async def lifespan(app: ASGIApp) -> AsyncIterator[None]:
        await state.startup()
        try:
            if original_lifespan:
                async with original_lifespan(app) as value:
                    yield value
            else:
                yield
        finally:
            if state._db is not None:
                await state._db.close()

    mcp_app.router.lifespan_context = lifespan

    if db is not None:
        return GroupAuthMiddleware(mcp_app, db, auth_store)
    return mcp_app


def main() -> None:
    server = create_server()
    server.run(transport="streamable-http", host="0.0.0.0", port=8765)


if __name__ == "__main__":
    main()
