from __future__ import annotations

import logging
import os
import re
import uuid
from dataclasses import dataclass
from typing import Any, Literal, Optional

import httpx
from mcp.server.fastmcp import FastMCP

logging.basicConfig(level=logging.INFO)
log = logging.getLogger("open-notebook-mcp")

mcp = FastMCP("open-notebook-mcp")

# -----------------------------
# Configuration and Constants
# -----------------------------

# Maximum number of items to return in list operations
MAX_LIMIT = 100

# Default timeout for HTTP requests (seconds)
DEFAULT_TIMEOUT_S = 30.0

def get_base_url() -> str:
    """Get the Open Notebook API base URL from environment."""
    return os.getenv("OPEN_NOTEBOOK_URL", "http://localhost:5055")

def get_auth_token() -> Optional[str]:
    """Get the authentication token from environment."""
    return os.getenv("OPEN_NOTEBOOK_PASSWORD")

def generate_request_id() -> str:
    """Generate a unique request ID for tracking."""
    return str(uuid.uuid4())

# -----------------------------
# Capability index (source of truth)
# -----------------------------
# Note: CAPABILITIES is a tuple to ensure immutability and prevent accidental modification.
# This is the single source of truth for all available tools in the server.

Detail = Literal["name", "summary", "full"]

@dataclass(frozen=True)
class Capability:
    name: str                 # tool name
    summary: str              # one-liner
    tags: tuple[str, ...]     # searchable tags
    args: dict[str, str]      # param -> type (stringified)
    returns: str              # return type (stringified)
    example: dict[str, Any]   # example call args
    typical_bytes: int        # typical response size (rough)

CAPABILITIES: tuple[Capability, ...] = (
    Capability(
        name="search_capabilities",
        summary="Search tools exposed by this server with progressive detail levels.",
        tags=("meta", "discovery", "progressive-disclosure"),
        args={"query": "str", "detail": "Literal['name','summary','full']", "limit": "int"},
        returns="dict[str, Any]",
        example={"query": "notebook", "detail": "summary", "limit": 10},
        typical_bytes=1200,
    ),
    # Notebooks API
    Capability(
        name="list_notebooks",
        summary="Get all notebooks with optional filtering and ordering.",
        tags=("notebooks", "list", "query"),
        args={"archived": "Optional[bool]", "order_by": "str", "limit": "int"},
        returns="dict[str, Any]",
        example={"archived": False, "order_by": "updated desc", "limit": 20},
        typical_bytes=2000,
    ),
    Capability(
        name="get_notebook",
        summary="Get a specific notebook by ID.",
        tags=("notebooks", "get", "read"),
        args={"notebook_id": "str"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123"},
        typical_bytes=500,
    ),
    Capability(
        name="create_notebook",
        summary="Create a new notebook.",
        tags=("notebooks", "create", "write"),
        args={"name": "str", "description": "Optional[str]"},
        returns="dict[str, Any]",
        example={"name": "My Research", "description": "AI research notebook"},
        typical_bytes=500,
    ),
    Capability(
        name="update_notebook",
        summary="Update a notebook.",
        tags=("notebooks", "update", "write"),
        args={"notebook_id": "str", "name": "Optional[str]", "description": "Optional[str]", "archived": "Optional[bool]"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "name": "Updated Name"},
        typical_bytes=500,
    ),
    Capability(
        name="delete_notebook",
        summary="Delete a notebook.",
        tags=("notebooks", "delete", "write"),
        args={"notebook_id": "str"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123"},
        typical_bytes=100,
    ),
    # Sources API
    Capability(
        name="list_sources",
        summary="Get all sources with optional filtering.",
        tags=("sources", "list", "query"),
        args={"notebook_id": "Optional[str]", "limit": "int", "offset": "int"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "limit": 20, "offset": 0},
        typical_bytes=3000,
    ),
    Capability(
        name="get_source",
        summary="Get a specific source by ID.",
        tags=("sources", "get", "read"),
        args={"source_id": "str"},
        returns="dict[str, Any]",
        example={"source_id": "source:abc123"},
        typical_bytes=2000,
    ),
    Capability(
        name="create_source",
        summary="Create a new source (link, upload, or text).",
        tags=("sources", "create", "write"),
        args={"notebook_id": "str", "type": "str", "url": "Optional[str]", "title": "Optional[str]", "embed": "bool"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "type": "link", "url": "https://example.com", "embed": True},
        typical_bytes=2000,
    ),
    Capability(
        name="update_source",
        summary="Update a source.",
        tags=("sources", "update", "write"),
        args={"source_id": "str", "title": "Optional[str]", "topics": "Optional[list[str]]"},
        returns="dict[str, Any]",
        example={"source_id": "source:abc123", "title": "New Title"},
        typical_bytes=2000,
    ),
    Capability(
        name="delete_source",
        summary="Delete a source.",
        tags=("sources", "delete", "write"),
        args={"source_id": "str"},
        returns="dict[str, Any]",
        example={"source_id": "source:abc123"},
        typical_bytes=100,
    ),
    # Notes API
    Capability(
        name="list_notes",
        summary="Get all notes with optional filtering.",
        tags=("notes", "list", "query"),
        args={"notebook_id": "Optional[str]", "limit": "int", "offset": "int"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "limit": 20, "offset": 0},
        typical_bytes=2000,
    ),
    Capability(
        name="get_note",
        summary="Get a specific note by ID.",
        tags=("notes", "get", "read"),
        args={"note_id": "str"},
        returns="dict[str, Any]",
        example={"note_id": "note:abc123"},
        typical_bytes=1500,
    ),
    Capability(
        name="create_note",
        summary="Create a new note.",
        tags=("notes", "create", "write"),
        args={"notebook_id": "str", "title": "str", "content": "str", "topics": "Optional[list[str]]"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "title": "My Note", "content": "Note content"},
        typical_bytes=1500,
    ),
    Capability(
        name="update_note",
        summary="Update a note.",
        tags=("notes", "update", "write"),
        args={"note_id": "str", "title": "Optional[str]", "content": "Optional[str]", "topics": "Optional[list[str]]"},
        returns="dict[str, Any]",
        example={"note_id": "note:abc123", "title": "Updated Title"},
        typical_bytes=1500,
    ),
    Capability(
        name="delete_note",
        summary="Delete a note.",
        tags=("notes", "delete", "write"),
        args={"note_id": "str"},
        returns="dict[str, Any]",
        example={"note_id": "note:abc123"},
        typical_bytes=100,
    ),
    # Search API
    Capability(
        name="search",
        summary="Search content using vector or text search.",
        tags=("search", "query", "vector"),
        args={"query": "str", "type": "str", "notebook_id": "Optional[str]", "limit": "int"},
        returns="dict[str, Any]",
        example={"query": "AI research", "type": "vector", "limit": 10},
        typical_bytes=3000,
    ),
    Capability(
        name="ask_question",
        summary="Ask a question about your content with detailed control.",
        tags=("search", "ask", "ai", "question"),
        args={"question": "str", "strategy_model": "str", "answer_model": "str", "final_answer_model": "str"},
        returns="dict[str, Any]",
        example={"question": "What are the main AI applications?", "strategy_model": "model:abc", "answer_model": "model:abc", "final_answer_model": "model:abc"},
        typical_bytes=5000,
    ),
    Capability(
        name="ask_simple",
        summary="Ask a question about your content with simplified interface.",
        tags=("search", "ask", "ai", "question", "simple"),
        args={"question": "str", "strategy_model": "str", "answer_model": "str", "final_answer_model": "str"},
        returns="dict[str, Any]",
        example={"question": "Summarize my AI research", "strategy_model": "model:abc", "answer_model": "model:abc", "final_answer_model": "model:abc"},
        typical_bytes=4000,
    ),
    # Models API
    Capability(
        name="list_models",
        summary="Get all configured AI models.",
        tags=("models", "list", "ai"),
        args={"limit": "int"},
        returns="dict[str, Any]",
        example={"limit": 50},
        typical_bytes=2000,
    ),
    Capability(
        name="get_model",
        summary="Get a specific model by ID.",
        tags=("models", "get", "read", "ai"),
        args={"model_id": "str"},
        returns="dict[str, Any]",
        example={"model_id": "model:abc123"},
        typical_bytes=500,
    ),
    Capability(
        name="create_model",
        summary="Create a new AI model configuration.",
        tags=("models", "create", "write", "ai"),
        args={"name": "str", "provider": "str", "type": "str"},
        returns="dict[str, Any]",
        example={"name": "gpt-4", "provider": "openai", "type": "language"},
        typical_bytes=500,
    ),
    Capability(
        name="delete_model",
        summary="Delete a model configuration.",
        tags=("models", "delete", "write", "ai"),
        args={"model_id": "str"},
        returns="dict[str, Any]",
        example={"model_id": "model:abc123"},
        typical_bytes=100,
    ),
    Capability(
        name="get_default_models",
        summary="Get default model configurations.",
        tags=("models", "defaults", "ai"),
        args={},
        returns="dict[str, Any]",
        example={},
        typical_bytes=1000,
    ),
    # Chat API
    Capability(
        name="list_chat_sessions",
        summary="Get all chat sessions with optional filtering.",
        tags=("chat", "sessions", "list"),
        args={"notebook_id": "Optional[str]", "limit": "int"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "limit": 20},
        typical_bytes=2000,
    ),
    Capability(
        name="create_chat_session",
        summary="Create a new chat session.",
        tags=("chat", "sessions", "create", "write"),
        args={"notebook_id": "str", "title": "str"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123", "title": "Research Discussion"},
        typical_bytes=500,
    ),
    Capability(
        name="get_chat_session",
        summary="Get a specific chat session by ID.",
        tags=("chat", "sessions", "get", "read"),
        args={"session_id": "str"},
        returns="dict[str, Any]",
        example={"session_id": "session:abc123"},
        typical_bytes=3000,
    ),
    Capability(
        name="update_chat_session",
        summary="Update a chat session.",
        tags=("chat", "sessions", "update", "write"),
        args={"session_id": "str", "title": "Optional[str]"},
        returns="dict[str, Any]",
        example={"session_id": "session:abc123", "title": "Updated Title"},
        typical_bytes=500,
    ),
    Capability(
        name="delete_chat_session",
        summary="Delete a chat session.",
        tags=("chat", "sessions", "delete", "write"),
        args={"session_id": "str"},
        returns="dict[str, Any]",
        example={"session_id": "session:abc123"},
        typical_bytes=100,
    ),
    Capability(
        name="execute_chat",
        summary="Send a message in a chat session.",
        tags=("chat", "execute", "message", "ai"),
        args={"session_id": "str", "message": "str", "context": "Optional[dict]"},
        returns="dict[str, Any]",
        example={"session_id": "session:abc123", "message": "What are the key insights?"},
        typical_bytes=3000,
    ),
    Capability(
        name="get_chat_context",
        summary="Build context for a chat conversation.",
        tags=("chat", "context", "build"),
        args={"notebook_id": "str", "context_config": "Optional[dict]"},
        returns="dict[str, Any]",
        example={"notebook_id": "notebook:abc123"},
        typical_bytes=5000,
    ),
    # Settings API
    Capability(
        name="get_settings",
        summary="Get application settings.",
        tags=("settings", "get", "config"),
        args={},
        returns="dict[str, Any]",
        example={},
        typical_bytes=1000,
    ),
    Capability(
        name="update_settings",
        summary="Update application settings.",
        tags=("settings", "update", "write", "config"),
        args={"settings": "dict"},
        returns="dict[str, Any]",
        example={"settings": {"theme": "dark"}},
        typical_bytes=1000,
    ),
)

def _normalize(s: str) -> str:
    return re.sub(r"[^a-z0-9]+", " ", s.lower()).strip()

def _match_score(q: str, cap: Capability) -> int:
    if not q:
        return 0
    qn = _normalize(q)
    hay = " ".join([cap.name, cap.summary, " ".join(cap.tags)])
    hn = _normalize(hay)
    score = 0
    for token in qn.split():
        if token in hn:
            score += 1
    return score

@mcp.tool()
def search_capabilities(query: str = "", detail: Detail = "summary", limit: int = 20) -> dict[str, Any]:
    """Search tools exposed by this server with progressive detail levels.
    
    Args:
        query: Search query to filter tools
        detail: Level of detail - 'name' (minimal), 'summary' (default), or 'full' (complete)
        limit: Maximum number of results (1-50)
    
    Returns:
        Dictionary with request_id, query, detail, count, matches, and hint fields
    """
    request_id = generate_request_id()
    limit = max(1, min(int(limit), 50))

    scored = [(c, _match_score(query, c)) for c in CAPABILITIES]
    if query.strip():
        scored = [x for x in scored if x[1] > 0]
        scored.sort(key=lambda t: (-t[1], t[0].name))
    else:
        scored.sort(key=lambda t: t[0].name)

    matches = []
    for cap, _score in scored[:limit]:
        if detail == "name":
            matches.append({"name": cap.name})
        elif detail == "summary":
            matches.append({"name": cap.name, "summary": cap.summary, "tags": list(cap.tags)})
        else:
            matches.append({
                "name": cap.name,
                "summary": cap.summary,
                "tags": list(cap.tags),
                "args": cap.args,
                "returns": cap.returns,
                "example": cap.example,
                "typical_bytes": cap.typical_bytes,
            })

    return {
        "request_id": request_id,
        "query": query,
        "detail": detail,
        "count": len(matches),
        "matches": matches,
        "hint": "Use detail='name' for minimal context; detail='full' only when implementing a call.",
    }

# -----------------------------
# HTTP helper functions
# -----------------------------
# HTTP helper functions
# -----------------------------

async def make_request(
    method: str,
    endpoint: str,
    *,
    json_data: Optional[dict[str, Any]] = None,
    params: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    """Make an HTTP request to the Open Notebook API.
    
    Args:
        method: HTTP method (GET, POST, PUT, DELETE)
        endpoint: API endpoint path (e.g., '/api/notebooks')
        json_data: Optional JSON body for POST/PUT requests
        params: Optional query parameters
    
    Returns:
        Response JSON as dictionary
    
    Raises:
        Exception: If request fails with formatted error message
    """
    base_url = get_base_url()
    url = f"{base_url}{endpoint}"
    
    headers = {"Content-Type": "application/json"}
    auth_token = get_auth_token()
    if auth_token:
        headers["Authorization"] = f"Bearer {auth_token}"
    
    async with httpx.AsyncClient(follow_redirects=True, timeout=DEFAULT_TIMEOUT_S) as client:
        try:
            if method == "GET":
                r = await client.get(url, headers=headers, params=params)
            elif method == "POST":
                r = await client.post(url, headers=headers, json=json_data, params=params)
            elif method == "PUT":
                r = await client.put(url, headers=headers, json=json_data, params=params)
            elif method == "DELETE":
                r = await client.delete(url, headers=headers, params=params)
            else:
                raise ValueError(f"Unsupported HTTP method: {method}")
            
            r.raise_for_status()
            
            # Handle empty responses
            if not r.text:
                return {"message": "Success"}
            
            try:
                return r.json()
            except (ValueError, Exception) as json_err:
                # If response is not valid JSON, return text content
                log.warning(f"Non-JSON response from {endpoint}: {json_err}")
                return {"message": "Success", "content": r.text}
            
        except httpx.HTTPError as e:
            error_msg = str(e)
            if hasattr(e, 'response') and e.response is not None:
                try:
                    error_detail = e.response.json()
                    error_msg = error_detail.get("detail", error_msg)
                except (ValueError, AttributeError):
                    # If error response is not JSON, use text or default message
                    error_msg = e.response.text or error_msg
            
            raise Exception(f"API request failed: {error_msg}")

# -----------------------------
# Notebooks API Tools
# -----------------------------

@mcp.tool()
async def list_notebooks(
    archived: Optional[bool] = None,
    order_by: str = "updated desc",
    limit: int = 20,
) -> dict[str, Any]:
    """Get all notebooks with optional filtering and ordering.
    
    Args:
        archived: Filter by archived status (None = all, True = archived only, False = active only)
        order_by: Order by field and direction (e.g., 'created desc', 'name asc')
        limit: Maximum number of results (1-100)
    
    Returns:
        Dictionary with notebooks list and metadata
    """
    limit = max(1, min(limit, MAX_LIMIT))
    params = {"order_by": order_by}
    if archived is not None:
        params["archived"] = archived
    
    notebooks = await make_request("GET", "/api/notebooks", params=params)
    
    # Limit results
    if isinstance(notebooks, list):
        notebooks = notebooks[:limit]
    
    return {
        "request_id": generate_request_id(),
        "count": len(notebooks) if isinstance(notebooks, list) else 0,
        "notebooks": notebooks,
    }

@mcp.tool()
async def get_notebook(notebook_id: str) -> dict[str, Any]:
    """Get a specific notebook by ID.
    
    Args:
        notebook_id: Notebook ID (e.g., 'notebook:abc123')
    
    Returns:
        Notebook details
    """
    notebook = await make_request("GET", f"/api/notebooks/{notebook_id}")
    return {
        "request_id": generate_request_id(),
        "notebook": notebook,
    }

@mcp.tool()
async def create_notebook(name: str, description: Optional[str] = None) -> dict[str, Any]:
    """Create a new notebook.
    
    Args:
        name: Notebook name
        description: Optional notebook description
    
    Returns:
        Created notebook details
    """
    data = {"name": name}
    if description is not None:
        data["description"] = description
    
    notebook = await make_request("POST", "/api/notebooks", json_data=data)
    return {
        "request_id": generate_request_id(),
        "notebook": notebook,
    }

@mcp.tool()
async def update_notebook(
    notebook_id: str,
    name: Optional[str] = None,
    description: Optional[str] = None,
    archived: Optional[bool] = None,
) -> dict[str, Any]:
    """Update a notebook.
    
    Args:
        notebook_id: Notebook ID
        name: Optional new name
        description: Optional new description
        archived: Optional archived status
    
    Returns:
        Updated notebook details
    """
    data = {}
    if name is not None:
        data["name"] = name
    if description is not None:
        data["description"] = description
    if archived is not None:
        data["archived"] = archived
    
    notebook = await make_request("PUT", f"/api/notebooks/{notebook_id}", json_data=data)
    return {
        "request_id": generate_request_id(),
        "notebook": notebook,
    }

@mcp.tool()
async def delete_notebook(notebook_id: str) -> dict[str, Any]:
    """Delete a notebook.
    
    Args:
        notebook_id: Notebook ID
    
    Returns:
        Success message
    """
    result = await make_request("DELETE", f"/api/notebooks/{notebook_id}")
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

# -----------------------------
# Sources API Tools
# -----------------------------

@mcp.tool()
async def list_sources(
    notebook_id: Optional[str] = None,
    limit: int = 20,
    offset: int = 0,
) -> dict[str, Any]:
    """Get all sources with optional filtering.
    
    Args:
        notebook_id: Optional notebook ID to filter by
        limit: Maximum number of results (1-100)
        offset: Pagination offset
    
    Returns:
        Dictionary with sources list and metadata
    """
    limit = max(1, min(limit, MAX_LIMIT))
    params = {"limit": limit, "offset": offset}
    if notebook_id is not None:
        params["notebook_id"] = notebook_id
    
    sources = await make_request("GET", "/api/sources", params=params)
    return {
        "request_id": generate_request_id(),
        "count": len(sources) if isinstance(sources, list) else 0,
        "sources": sources,
    }

@mcp.tool()
async def get_source(source_id: str) -> dict[str, Any]:
    """Get a specific source by ID.
    
    Args:
        source_id: Source ID (e.g., 'source:abc123')
    
    Returns:
        Source details
    """
    source = await make_request("GET", f"/api/sources/{source_id}")
    return {
        "request_id": generate_request_id(),
        "source": source,
    }

@mcp.tool()
async def create_source(
    notebook_id: str,
    type: str,
    url: Optional[str] = None,
    title: Optional[str] = None,
    embed: bool = True,
) -> dict[str, Any]:
    """Create a new source (link, upload, or text).
    
    Args:
        notebook_id: Notebook ID to add source to
        type: Source type ('link', 'upload', or 'text')
        url: URL for link type sources
        title: Optional title
        embed: Whether to generate embeddings (default: True)
    
    Returns:
        Created source details
    """
    data = {
        "notebook_id": notebook_id,
        "type": type,
        "embed": embed,
    }
    if url is not None:
        data["url"] = url
    if title is not None:
        data["title"] = title
    
    source = await make_request("POST", "/api/sources", json_data=data)
    return {
        "request_id": generate_request_id(),
        "source": source,
    }

@mcp.tool()
async def update_source(
    source_id: str,
    title: Optional[str] = None,
    topics: Optional[list[str]] = None,
) -> dict[str, Any]:
    """Update a source.
    
    Args:
        source_id: Source ID
        title: Optional new title
        topics: Optional list of topics
    
    Returns:
        Updated source details
    """
    data = {}
    if title is not None:
        data["title"] = title
    if topics is not None:
        data["topics"] = topics
    
    source = await make_request("PUT", f"/api/sources/{source_id}", json_data=data)
    return {
        "request_id": generate_request_id(),
        "source": source,
    }

@mcp.tool()
async def delete_source(source_id: str) -> dict[str, Any]:
    """Delete a source.
    
    Args:
        source_id: Source ID
    
    Returns:
        Success message
    """
    result = await make_request("DELETE", f"/api/sources/{source_id}")
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

# -----------------------------
# Notes API Tools
# -----------------------------

@mcp.tool()
async def list_notes(
    notebook_id: Optional[str] = None,
    limit: int = 20,
    offset: int = 0,
) -> dict[str, Any]:
    """Get all notes with optional filtering.
    
    Args:
        notebook_id: Optional notebook ID to filter by
        limit: Maximum number of results (1-100)
        offset: Pagination offset
    
    Returns:
        Dictionary with notes list and metadata
    """
    limit = max(1, min(limit, MAX_LIMIT))
    params = {"limit": limit, "offset": offset}
    if notebook_id is not None:
        params["notebook_id"] = notebook_id
    
    notes = await make_request("GET", "/api/notes", params=params)
    return {
        "request_id": generate_request_id(),
        "count": len(notes) if isinstance(notes, list) else 0,
        "notes": notes,
    }

@mcp.tool()
async def get_note(note_id: str) -> dict[str, Any]:
    """Get a specific note by ID.
    
    Args:
        note_id: Note ID (e.g., 'note:abc123')
    
    Returns:
        Note details
    """
    note = await make_request("GET", f"/api/notes/{note_id}")
    return {
        "request_id": generate_request_id(),
        "note": note,
    }

@mcp.tool()
async def create_note(
    notebook_id: str,
    title: str,
    content: str,
    topics: Optional[list[str]] = None,
) -> dict[str, Any]:
    """Create a new note.
    
    Args:
        notebook_id: Notebook ID to add note to
        title: Note title
        content: Note content
        topics: Optional list of topics
    
    Returns:
        Created note details
    """
    data = {
        "notebook_id": notebook_id,
        "title": title,
        "content": content,
    }
    if topics is not None:
        data["topics"] = topics
    
    note = await make_request("POST", "/api/notes", json_data=data)
    return {
        "request_id": generate_request_id(),
        "note": note,
    }

@mcp.tool()
async def update_note(
    note_id: str,
    title: Optional[str] = None,
    content: Optional[str] = None,
    topics: Optional[list[str]] = None,
) -> dict[str, Any]:
    """Update a note.
    
    Args:
        note_id: Note ID
        title: Optional new title
        content: Optional new content
        topics: Optional list of topics
    
    Returns:
        Updated note details
    """
    data = {}
    if title is not None:
        data["title"] = title
    if content is not None:
        data["content"] = content
    if topics is not None:
        data["topics"] = topics
    
    note = await make_request("PUT", f"/api/notes/{note_id}", json_data=data)
    return {
        "request_id": generate_request_id(),
        "note": note,
    }

@mcp.tool()
async def delete_note(note_id: str) -> dict[str, Any]:
    """Delete a note.
    
    Args:
        note_id: Note ID
    
    Returns:
        Success message
    """
    result = await make_request("DELETE", f"/api/notes/{note_id}")
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

# -----------------------------
# Search API Tools
# -----------------------------

@mcp.tool()
async def search(
    query: str,
    type: str = "vector",
    notebook_id: Optional[str] = None,
    limit: int = 10,
) -> dict[str, Any]:
    """Search content using vector or text search.
    
    Args:
        query: Search query
        type: Search type ('vector' or 'text')
        notebook_id: Optional notebook ID to limit search
        limit: Maximum number of results (1-50)
    
    Returns:
        Search results
    """
    limit = max(1, min(limit, 50))
    data = {
        "query": query,
        "type": type,
        "limit": limit,
    }
    if notebook_id is not None:
        data["notebook_id"] = notebook_id
    
    results = await make_request("POST", "/api/search", json_data=data)
    return {
        "request_id": generate_request_id(),
        "results": results,
    }

@mcp.tool()
async def ask_question(
    question: str,
    strategy_model: str,
    answer_model: str,
    final_answer_model: str,
    notebook_id: Optional[str] = None,
) -> dict[str, Any]:
    """Ask a question about your content with detailed control.
    
    Args:
        question: Question to ask
        strategy_model: Model ID for strategy generation
        answer_model: Model ID for answering
        final_answer_model: Model ID for final answer synthesis
        notebook_id: Optional notebook ID to limit context
    
    Returns:
        Answer with sources and reasoning
    """
    data = {
        "question": question,
        "strategy_model": strategy_model,
        "answer_model": answer_model,
        "final_answer_model": final_answer_model,
    }
    if notebook_id is not None:
        data["notebook_id"] = notebook_id
    
    result = await make_request("POST", "/api/search/ask", json_data=data)
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

@mcp.tool()
async def ask_simple(
    question: str,
    strategy_model: str,
    answer_model: str,
    final_answer_model: str,
    notebook_id: Optional[str] = None,
) -> dict[str, Any]:
    """Ask a question about your content with simplified interface.
    
    Args:
        question: Question to ask
        strategy_model: Model ID for strategy generation
        answer_model: Model ID for answering
        final_answer_model: Model ID for final answer synthesis
        notebook_id: Optional notebook ID to limit context
    
    Returns:
        Simple answer
    """
    data = {
        "question": question,
        "strategy_model": strategy_model,
        "answer_model": answer_model,
        "final_answer_model": final_answer_model,
    }
    if notebook_id is not None:
        data["notebook_id"] = notebook_id
    
    result = await make_request("POST", "/api/search/ask/simple", json_data=data)
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

# -----------------------------
# Models API Tools
# -----------------------------

@mcp.tool()
async def list_models(limit: int = 50) -> dict[str, Any]:
    """Get all configured AI models.
    
    Args:
        limit: Maximum number of results (1-100)
    
    Returns:
        Dictionary with models list and metadata
    """
    limit = max(1, min(limit, MAX_LIMIT))
    models = await make_request("GET", "/api/models")
    
    if isinstance(models, list):
        models = models[:limit]
    
    return {
        "request_id": generate_request_id(),
        "count": len(models) if isinstance(models, list) else 0,
        "models": models,
    }

@mcp.tool()
async def get_model(model_id: str) -> dict[str, Any]:
    """Get a specific model by ID.
    
    Args:
        model_id: Model ID (e.g., 'model:abc123')
    
    Returns:
        Model details
    """
    model = await make_request("GET", f"/api/models/{model_id}")
    return {
        "request_id": generate_request_id(),
        "model": model,
    }

@mcp.tool()
async def create_model(name: str, provider: str, type: str) -> dict[str, Any]:
    """Create a new AI model configuration.
    
    Args:
        name: Model name (e.g., 'gpt-4', 'claude-3-opus')
        provider: Provider name (e.g., 'openai', 'anthropic')
        type: Model type (e.g., 'language', 'embedding')
    
    Returns:
        Created model details
    """
    data = {
        "name": name,
        "provider": provider,
        "type": type,
    }
    
    model = await make_request("POST", "/api/models", json_data=data)
    return {
        "request_id": generate_request_id(),
        "model": model,
    }

@mcp.tool()
async def delete_model(model_id: str) -> dict[str, Any]:
    """Delete a model configuration.
    
    Args:
        model_id: Model ID
    
    Returns:
        Success message
    """
    result = await make_request("DELETE", f"/api/models/{model_id}")
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

@mcp.tool()
async def get_default_models() -> dict[str, Any]:
    """Get default model configurations.
    
    Returns:
        Default models configuration
    """
    defaults = await make_request("GET", "/api/models/defaults")
    return {
        "request_id": generate_request_id(),
        "defaults": defaults,
    }

# -----------------------------
# Chat API Tools
# -----------------------------

@mcp.tool()
async def list_chat_sessions(
    notebook_id: Optional[str] = None,
    limit: int = 20,
) -> dict[str, Any]:
    """Get all chat sessions with optional filtering.
    
    Args:
        notebook_id: Optional notebook ID to filter by
        limit: Maximum number of results (1-100)
    
    Returns:
        Dictionary with sessions list and metadata
    """
    limit = max(1, min(limit, MAX_LIMIT))
    params = {}
    if notebook_id is not None:
        params["notebook_id"] = notebook_id
    
    sessions = await make_request("GET", "/api/chat/sessions", params=params)
    
    if isinstance(sessions, list):
        sessions = sessions[:limit]
    
    return {
        "request_id": generate_request_id(),
        "count": len(sessions) if isinstance(sessions, list) else 0,
        "sessions": sessions,
    }

@mcp.tool()
async def create_chat_session(notebook_id: str, title: str) -> dict[str, Any]:
    """Create a new chat session.
    
    Args:
        notebook_id: Notebook ID for the session
        title: Session title
    
    Returns:
        Created session details
    """
    data = {
        "notebook_id": notebook_id,
        "title": title,
    }
    
    session = await make_request("POST", "/api/chat/sessions", json_data=data)
    return {
        "request_id": generate_request_id(),
        "session": session,
    }

@mcp.tool()
async def get_chat_session(session_id: str) -> dict[str, Any]:
    """Get a specific chat session by ID.
    
    Args:
        session_id: Session ID (e.g., 'session:abc123')
    
    Returns:
        Session details with message history
    """
    session = await make_request("GET", f"/api/chat/sessions/{session_id}")
    return {
        "request_id": generate_request_id(),
        "session": session,
    }

@mcp.tool()
async def update_chat_session(
    session_id: str,
    title: Optional[str] = None,
) -> dict[str, Any]:
    """Update a chat session.
    
    Args:
        session_id: Session ID
        title: Optional new title
    
    Returns:
        Updated session details
    """
    data = {}
    if title is not None:
        data["title"] = title
    
    session = await make_request("PUT", f"/api/chat/sessions/{session_id}", json_data=data)
    return {
        "request_id": generate_request_id(),
        "session": session,
    }

@mcp.tool()
async def delete_chat_session(session_id: str) -> dict[str, Any]:
    """Delete a chat session.
    
    Args:
        session_id: Session ID
    
    Returns:
        Success message
    """
    result = await make_request("DELETE", f"/api/chat/sessions/{session_id}")
    return {
        "request_id": generate_request_id(),
        "result": result,
    }

@mcp.tool()
async def execute_chat(
    session_id: str,
    message: str,
    context: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    """Send a message in a chat session.
    
    Args:
        session_id: Session ID
        message: Message to send
        context: Optional context data for the conversation
    
    Returns:
        Chat response with AI message
    """
    data = {
        "session_id": session_id,
        "message": message,
    }
    if context is not None:
        data["context"] = context
    
    response = await make_request("POST", "/api/chat/execute", json_data=data)
    return {
        "request_id": generate_request_id(),
        "response": response,
    }

@mcp.tool()
async def get_chat_context(
    notebook_id: str,
    context_config: Optional[dict[str, Any]] = None,
) -> dict[str, Any]:
    """Build context for a chat conversation.
    
    Args:
        notebook_id: Notebook ID
        context_config: Optional context configuration
    
    Returns:
        Built context data
    """
    data = {
        "notebook_id": notebook_id,
    }
    if context_config is not None:
        data["context_config"] = context_config
    
    context = await make_request("POST", "/api/chat/context", json_data=data)
    return {
        "request_id": generate_request_id(),
        "context": context,
    }

# -----------------------------
# Settings API Tools
# -----------------------------

@mcp.tool()
async def get_settings() -> dict[str, Any]:
    """Get application settings.
    
    Returns:
        Application settings
    """
    settings = await make_request("GET", "/api/settings")
    return {
        "request_id": generate_request_id(),
        "settings": settings,
    }

@mcp.tool()
async def update_settings(settings: dict[str, Any]) -> dict[str, Any]:
    """Update application settings.
    
    Args:
        settings: Settings dictionary to update
    
    Returns:
        Updated settings
    """
    result = await make_request("PUT", "/api/settings", json_data=settings)
    return {
        "request_id": generate_request_id(),
        "settings": result,
    }

# -----------------------------
# Dual transport entrypoint
# -----------------------------

def main() -> None:
    """Main entry point for the MCP server."""
    # Default: stdio (Codex/local)
    transport = os.getenv("MCP_TRANSPORT", "stdio").strip().lower()

    if transport == "stdio":
        mcp.run(transport="stdio")
        return

    # Streamable HTTP (prod/remote)
    host = os.getenv("HOST", "127.0.0.1")
    port = int(os.getenv("PORT", "8000"))
    path = os.getenv("MCP_PATH", "/mcp")
    stateless_http = os.getenv("STATELESS_HTTP", "1") == "1"
    json_response = os.getenv("JSON_RESPONSE", "1") == "1"

    mcp.run(
        transport="streamable-http",
        host=host,
        port=port,
        path=path,
        stateless_http=stateless_http,
        json_response=json_response,
    )

if __name__ == "__main__":
    main()
