main.py

"""
OneNote Integration App for Omi

This app provides OneNote integration through Microsoft OAuth2 authentication
and chat tools for creating, searching, and managing notes.
"""
import os
import sys
import secrets
import difflib
import urllib.parse
from datetime import datetime, timedelta
from typing import Optional, Dict, Any, List, Tuple
import html
import re

import requests
from dotenv import load_dotenv
from fastapi import FastAPI, HTTPException, Request, Query
from fastapi.responses import HTMLResponse, RedirectResponse, JSONResponse
from fastapi.staticfiles import StaticFiles
from fastapi.templating import Jinja2Templates

from db import (
    store_onenote_tokens,
    get_onenote_tokens,
    delete_onenote_tokens,
    update_onenote_tokens,
    store_oauth_state,
    get_oauth_state,
    delete_oauth_state,
    get_user_settings,
)
from models import (
    ChatToolResponse,
    OneNoteNotebook,
    OneNoteSection,
    OneNotePage,
    OneNoteUser,
)


def log(msg: str):
    """Print and flush immediately for Railway logging."""
    print(msg)
    sys.stdout.flush()


load_dotenv()

# Microsoft Azure AD / Graph API Configuration
MICROSOFT_CLIENT_ID = os.getenv("MICROSOFT_CLIENT_ID", "")
MICROSOFT_CLIENT_SECRET = os.getenv("MICROSOFT_CLIENT_SECRET", "")
MICROSOFT_REDIRECT_URI = os.getenv("MICROSOFT_REDIRECT_URI", "http://localhost:8080/auth/microsoft/callback")
MICROSOFT_TENANT = os.getenv("MICROSOFT_TENANT", "common")  # 'common' for multi-tenant

# Microsoft OAuth2 endpoints
MICROSOFT_AUTH_URL = f"https://login.microsoftonline.com/{MICROSOFT_TENANT}/oauth2/v2.0/authorize"
MICROSOFT_TOKEN_URL = f"https://login.microsoftonline.com/{MICROSOFT_TENANT}/oauth2/v2.0/token"
GRAPH_API_BASE = "https://graph.microsoft.com/v1.0"

# Required scopes for OneNote access
MICROSOFT_SCOPES = [
    "User.Read",
    "Notes.Read",
    "Notes.ReadWrite",
    "Notes.Create",
    "offline_access",  # For refresh tokens
]

app = FastAPI(
    title="OneNote Omi Integration",
    description="OneNote integration for Omi - Create and search notes using voice",
    version="1.0.0"
)

# Mount static files and templates
templates_dir = os.path.join(os.path.dirname(__file__), "templates")
if os.path.exists(templates_dir):
    static_dir = os.path.join(templates_dir, "static")
    if os.path.exists(static_dir):
        app.mount("/static", StaticFiles(directory=static_dir), name="static")
templates = Jinja2Templates(directory=templates_dir)


# ============================================
# Helper Functions
# ============================================

def get_valid_access_token(uid: str) -> Optional[str]:
    """Get a valid access token, refreshing if necessary."""
    tokens = get_onenote_tokens(uid)
    if not tokens:
        return None

    access_token = tokens.get("access_token")
    refresh_token = tokens.get("refresh_token")
    expires_at_str = tokens.get("expires_at")

    if not access_token:
        return None

    # Check if token is expired
    if expires_at_str:
        try:
            expires_at = datetime.fromisoformat(expires_at_str)
            # Refresh if expires in less than 5 minutes
            if datetime.utcnow() >= expires_at - timedelta(minutes=5):
                log(f"Token expired or expiring soon for {uid}, refreshing...")
                if refresh_token:
                    new_tokens = refresh_access_token(refresh_token)
                    if new_tokens:
                        access_token = new_tokens["access_token"]
                        new_refresh = new_tokens.get("refresh_token", refresh_token)
                        new_expires = datetime.utcnow() + timedelta(seconds=new_tokens.get("expires_in", 3600))
                        update_onenote_tokens(uid, access_token, new_refresh, new_expires)
                        log(f"Token refreshed successfully for {uid}")
                    else:
                        log(f"Failed to refresh token for {uid}")
                        return None
        except Exception as e:
            log(f"Error checking token expiry: {e}")

    return access_token


def refresh_access_token(refresh_token: str) -> Optional[Dict[str, Any]]:
    """Refresh the access token using the refresh token."""
    try:
        data = {
            "client_id": MICROSOFT_CLIENT_ID,
            "client_secret": MICROSOFT_CLIENT_SECRET,
            "refresh_token": refresh_token,
            "grant_type": "refresh_token",
            "scope": " ".join(MICROSOFT_SCOPES),
        }

        response = requests.post(MICROSOFT_TOKEN_URL, data=data)
        if response.status_code == 200:
            return response.json()
        else:
            log(f"Token refresh failed: {response.status_code} - {response.text}")
            return None
    except Exception as e:
        log(f"Token refresh error: {e}")
        return None


def graph_api_request(
    access_token: str,
    endpoint: str,
    method: str = "GET",
    json_data: Optional[Dict] = None,
    content_type: str = "application/json"
) -> Tuple[Optional[Dict], Optional[str]]:
    """Make a request to Microsoft Graph API."""
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": content_type,
    }

    url = f"{GRAPH_API_BASE}{endpoint}"

    try:
        if method == "GET":
            response = requests.get(url, headers=headers)
        elif method == "POST":
            if content_type == "application/json":
                response = requests.post(url, headers=headers, json=json_data)
            else:
                response = requests.post(url, headers=headers, data=json_data)
        elif method == "PATCH":
            response = requests.patch(url, headers=headers, json=json_data)
        elif method == "DELETE":
            response = requests.delete(url, headers=headers)
        else:
            return None, f"Unsupported method: {method}"

        if response.status_code in [200, 201, 204]:
            if response.status_code == 204:
                return {}, None
            return response.json(), None
        else:
            error_msg = f"API error {response.status_code}: {response.text}"
            log(error_msg)
            return None, error_msg

    except Exception as e:
        error_msg = f"Request failed: {str(e)}"
        log(error_msg)
        return None, error_msg


def get_current_user(uid: str) -> Optional[OneNoteUser]:
    """Get the current Microsoft user info."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return None

    data, error = graph_api_request(access_token, "/me")
    if error or not data:
        return None

    return OneNoteUser(
        id=data.get("id", ""),
        display_name=data.get("displayName", ""),
        email=data.get("mail") or data.get("userPrincipalName"),
        user_principal_name=data.get("userPrincipalName")
    )


def get_notebooks(uid: str) -> List[OneNoteNotebook]:
    """Get the user's notebooks."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return []

    data, error = graph_api_request(access_token, "/me/onenote/notebooks")
    if error or not data:
        return []

    notebooks = []
    for nb in data.get("value", []):
        notebooks.append(OneNoteNotebook(
            id=nb.get("id", ""),
            display_name=nb.get("displayName", ""),
            created_datetime=nb.get("createdDateTime"),
            last_modified_datetime=nb.get("lastModifiedDateTime"),
            is_default=nb.get("isDefault", False),
            is_shared=nb.get("isShared", False)
        ))
    return notebooks


def get_sections(uid: str, notebook_id: Optional[str] = None) -> List[OneNoteSection]:
    """Get sections, optionally filtered by notebook."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return []

    if notebook_id:
        endpoint = f"/me/onenote/notebooks/{notebook_id}/sections"
    else:
        endpoint = "/me/onenote/sections"

    data, error = graph_api_request(access_token, endpoint)
    if error or not data:
        return []

    sections = []
    for s in data.get("value", []):
        sections.append(OneNoteSection(
            id=s.get("id", ""),
            display_name=s.get("displayName", ""),
            notebook_id=notebook_id,
            created_datetime=s.get("createdDateTime"),
            last_modified_datetime=s.get("lastModifiedDateTime")
        ))
    return sections


def get_pages(uid: str, section_id: Optional[str] = None, search_query: Optional[str] = None, limit: int = 20) -> List[OneNotePage]:
    """Get pages, optionally filtered by section or search query."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return []

    if section_id:
        endpoint = f"/me/onenote/sections/{section_id}/pages?$top={limit}&$orderby=lastModifiedDateTime desc"
    elif search_query:
        # Use search endpoint
        encoded_query = urllib.parse.quote(search_query)
        endpoint = f"/me/onenote/pages?$search={encoded_query}&$top={limit}"
    else:
        endpoint = f"/me/onenote/pages?$top={limit}&$orderby=lastModifiedDateTime desc"

    data, error = graph_api_request(access_token, endpoint)

    # If search returned no results, try fetching all pages and filter locally
    if search_query and (error or not data or not data.get("value")):
        log(f"Graph search returned no results, trying local search for '{search_query}'")
        all_pages_endpoint = f"/me/onenote/pages?$top=50&$orderby=lastModifiedDateTime desc"
        data, error = graph_api_request(access_token, all_pages_endpoint)

        if data and data.get("value"):
            query_lower = search_query.lower()
            filtered = []
            for p in data.get("value", []):
                title = (p.get("title") or "").lower()
                preview = (p.get("contentPreviewText") or "").lower()
                if query_lower in title or query_lower in preview:
                    filtered.append(p)
            data["value"] = filtered[:limit]
            log(f"Local search found {len(filtered)} matches")

    if error or not data:
        return []

    pages = []
    for p in data.get("value", []):
        pages.append(OneNotePage(
            id=p.get("id", ""),
            title=p.get("title", "Untitled"),
            section_id=section_id,
            created_datetime=p.get("createdDateTime"),
            last_modified_datetime=p.get("lastModifiedDateTime"),
            content_url=p.get("contentUrl"),
            preview_text=p.get("contentPreviewText")
        ))
    return pages


def get_page_content(uid: str, page_id: str) -> Optional[str]:
    """Get the HTML content of a page."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return None

    url = f"{GRAPH_API_BASE}/me/onenote/pages/{page_id}/content"
    headers = {"Authorization": f"Bearer {access_token}"}

    try:
        response = requests.get(url, headers=headers)
        if response.status_code == 200:
            return response.text
        else:
            log(f"Failed to get page content: {response.status_code}")
            return None
    except Exception as e:
        log(f"Error getting page content: {e}")
        return None


def html_to_text(html_content: str) -> str:
    """Convert HTML content to plain text."""
    # Remove script and style elements
    text = re.sub(r'<script[^>]*>.*?</script>', '', html_content, flags=re.DOTALL | re.IGNORECASE)
    text = re.sub(r'<style[^>]*>.*?</style>', '', text, flags=re.DOTALL | re.IGNORECASE)
    # Replace common block elements with newlines
    text = re.sub(r'<br\s*/?>', '\n', text, flags=re.IGNORECASE)
    text = re.sub(r'</(p|div|h[1-6]|li|tr)>', '\n', text, flags=re.IGNORECASE)
    # Remove all remaining tags
    text = re.sub(r'<[^>]+>', '', text)
    # Decode HTML entities
    text = html.unescape(text)
    # Clean up whitespace
    text = re.sub(r'\n\s*\n', '\n\n', text)
    text = text.strip()
    return text


def fuzzy_match_notebook(name: str, notebooks: List[OneNoteNotebook], threshold: float = 0.5) -> Tuple[Optional[OneNoteNotebook], float]:
    """Fuzzy match a notebook name."""
    if not notebooks:
        return None, 0.0

    name_lower = name.lower().strip()
    best_match = None
    best_score = 0.0

    for nb in notebooks:
        nb_name_lower = nb.display_name.lower()
        score = difflib.SequenceMatcher(None, name_lower, nb_name_lower).ratio()

        # Boost if substring
        if name_lower in nb_name_lower or nb_name_lower.startswith(name_lower):
            score = max(score, 0.85)

        if score > best_score:
            best_score = score
            best_match = nb

    if best_score >= threshold:
        return best_match, best_score
    return None, best_score


def fuzzy_match_section(name: str, sections: List[OneNoteSection], threshold: float = 0.5) -> Tuple[Optional[OneNoteSection], float]:
    """Fuzzy match a section name."""
    if not sections:
        return None, 0.0

    name_lower = name.lower().strip()
    best_match = None
    best_score = 0.0

    for s in sections:
        s_name_lower = s.display_name.lower()
        score = difflib.SequenceMatcher(None, name_lower, s_name_lower).ratio()

        # Boost if substring
        if name_lower in s_name_lower or s_name_lower.startswith(name_lower):
            score = max(score, 0.85)

        if score > best_score:
            best_score = score
            best_match = s

    if best_score >= threshold:
        return best_match, best_score
    return None, best_score


def create_page(uid: str, section_id: str, title: str, content: str) -> Tuple[Optional[OneNotePage], Optional[str]]:
    """Create a new page in a section."""
    access_token = get_valid_access_token(uid)
    if not access_token:
        return None, "Not authenticated"

    # OneNote pages are created with HTML content
    # The content must include proper HTML structure
    html_content = f"""<!DOCTYPE html>
<html>
<head>
    <title>{html.escape(title)}</title>
</head>
<body>
    <h1>{html.escape(title)}</h1>
    <p>{html.escape(content).replace(chr(10), '</p><p>')}</p>
</body>
</html>"""

    url = f"{GRAPH_API_BASE}/me/onenote/sections/{section_id}/pages"
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/xhtml+xml",
    }

    try:
        response = requests.post(url, headers=headers, data=html_content.encode('utf-8'))
        if response.status_code in [200, 201]:
            data = response.json()
            page = OneNotePage(
                id=data.get("id", ""),
                title=data.get("title", title),
                section_id=section_id,
                created_datetime=data.get("createdDateTime"),
                last_modified_datetime=data.get("lastModifiedDateTime")
            )
            return page, None
        else:
            error_msg = f"Failed to create page: {response.status_code} - {response.text}"
            log(error_msg)
            return None, error_msg
    except Exception as e:
        error_msg = f"Error creating page: {str(e)}"
        log(error_msg)
        return None, error_msg


def get_default_section(uid: str) -> Optional[OneNoteSection]:
    """Get the default section (first section of default notebook, or first available)."""
    notebooks = get_notebooks(uid)
    if not notebooks:
        return None

    # Try to find default notebook first
    default_nb = next((nb for nb in notebooks if nb.is_default), None)
    if not default_nb:
        default_nb = notebooks[0]

    sections = get_sections(uid, default_nb.id)
    if sections:
        return sections[0]

    return None


# ============================================
# OAuth Endpoints
# ============================================

@app.get("/", response_class=HTMLResponse)
async def home(request: Request, uid: Optional[str] = None):
    """Home page / App settings page."""
    if not uid:
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": "Missing user ID"
        })

    tokens = get_onenote_tokens(uid)
    authenticated = tokens is not None

    user_info = None
    notebooks = []
    if authenticated:
        user_info = get_current_user(uid)
        if user_info:
            notebooks = get_notebooks(uid)

    return templates.TemplateResponse("setup.html", {
        "request": request,
        "uid": uid,
        "authenticated": authenticated,
        "user_info": user_info,
        "notebooks": notebooks,
    })


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {"status": "healthy", "service": "onenote-omi"}


@app.get("/auth/microsoft")
async def microsoft_auth(uid: str):
    """Initiate Microsoft OAuth2 flow."""
    if not uid:
        raise HTTPException(status_code=400, detail="User ID is required")

    if not MICROSOFT_CLIENT_ID or not MICROSOFT_CLIENT_SECRET:
        raise HTTPException(status_code=500, detail="Microsoft credentials not configured")

    # Generate state for CSRF protection
    state = f"{uid}:{secrets.token_urlsafe(32)}"
    store_oauth_state(uid, state)

    # Build authorization URL
    params = {
        "client_id": MICROSOFT_CLIENT_ID,
        "response_type": "code",
        "redirect_uri": MICROSOFT_REDIRECT_URI,
        "scope": " ".join(MICROSOFT_SCOPES),
        "state": state,
        "response_mode": "query",
    }

    auth_url = f"{MICROSOFT_AUTH_URL}?{urllib.parse.urlencode(params)}"
    return RedirectResponse(url=auth_url)


@app.get("/auth/microsoft/callback")
async def microsoft_callback(
    request: Request,
    code: Optional[str] = None,
    state: Optional[str] = None,
    error: Optional[str] = None,
    error_description: Optional[str] = None
):
    """Handle Microsoft OAuth2 callback."""
    if error:
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": f"Authorization failed: {error_description or error}"
        })

    if not code or not state:
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": "Invalid callback parameters"
        })

    # Extract uid from state
    try:
        uid, _ = state.split(":", 1)
    except ValueError:
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": "Invalid state parameter"
        })

    # Verify state matches what we stored
    stored_state = get_oauth_state(uid)
    if stored_state != state:
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": "State mismatch - possible CSRF attack"
        })

    # Clean up state
    delete_oauth_state(uid)

    # Exchange code for tokens
    try:
        token_data = {
            "client_id": MICROSOFT_CLIENT_ID,
            "client_secret": MICROSOFT_CLIENT_SECRET,
            "code": code,
            "redirect_uri": MICROSOFT_REDIRECT_URI,
            "grant_type": "authorization_code",
            "scope": " ".join(MICROSOFT_SCOPES),
        }

        response = requests.post(MICROSOFT_TOKEN_URL, data=token_data)

        if response.status_code != 200:
            error_data = response.json()
            error_msg = error_data.get("error_description", error_data.get("error", "Unknown error"))
            log(f"Token exchange failed: {error_msg}")
            return templates.TemplateResponse("setup.html", {
                "request": request,
                "authenticated": False,
                "error": f"Failed to get access token: {error_msg}"
            })

        tokens = response.json()
        access_token = tokens["access_token"]
        refresh_token = tokens.get("refresh_token", "")
        expires_in = tokens.get("expires_in", 3600)
        expires_at = datetime.utcnow() + timedelta(seconds=expires_in)

        # Store tokens
        store_onenote_tokens(
            uid,
            access_token,
            refresh_token,
            expires_at,
            tokens.get("token_type", "Bearer")
        )

        # Redirect to home with uid
        return RedirectResponse(url=f"/?uid={uid}")

    except Exception as e:
        log(f"OAuth error: {e}")
        return templates.TemplateResponse("setup.html", {
            "request": request,
            "authenticated": False,
            "error": f"Failed to complete authentication: {str(e)}"
        })


@app.get("/setup/onenote", tags=["setup"])
async def check_setup(uid: str):
    """Check if the user has completed OneNote setup (used by Omi)."""
    tokens = get_onenote_tokens(uid)
    return {"is_setup_completed": tokens is not None}


@app.get("/disconnect")
async def disconnect_onenote(uid: str):
    """Disconnect OneNote account."""
    delete_onenote_tokens(uid)
    return RedirectResponse(url=f"/?uid={uid}")


# ============================================
# Chat Tool Endpoints
# ============================================

@app.post("/tools/create_note", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_create_note(request: Request):
    """
    Create a new OneNote page.
    Chat tool for Omi - creates a note with the specified content.
    """
    try:
        body = await request.json()
        log(f"=== CREATE_NOTE START ===")
        log(f"Request: {body}")

        uid = body.get("uid")
        title = body.get("title", "Untitled Note")
        content = body.get("content", "")
        notebook_name = body.get("notebook")
        section_name = body.get("section")

        if not uid:
            return ChatToolResponse(error="User ID is required")

        if not content:
            return ChatToolResponse(error="Note content is required")

        # Check authentication
        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        # Find the target section
        section = None
        notebook_info = None

        if notebook_name:
            notebooks = get_notebooks(uid)
            notebook_match, _ = fuzzy_match_notebook(notebook_name, notebooks)
            if notebook_match:
                notebook_info = notebook_match
                sections = get_sections(uid, notebook_match.id)

                if section_name:
                    section_match, _ = fuzzy_match_section(section_name, sections)
                    if section_match:
                        section = section_match
                    else:
                        section_names = [s.display_name for s in sections[:5]]
                        return ChatToolResponse(
                            error=f"Could not find section '{section_name}' in notebook '{notebook_match.display_name}'. Available sections: {', '.join(section_names)}"
                        )
                elif sections:
                    section = sections[0]  # Use first section
            else:
                notebook_names = [nb.display_name for nb in notebooks[:5]]
                return ChatToolResponse(
                    error=f"Could not find notebook '{notebook_name}'. Your notebooks: {', '.join(notebook_names)}"
                )
        else:
            # Use default section
            section = get_default_section(uid)

        if not section:
            return ChatToolResponse(
                error="Could not find a section to create the note in. Please create a notebook and section in OneNote first."
            )

        # Create the page
        page, error = create_page(uid, section.id, title, content)
        if error:
            return ChatToolResponse(error=error)

        # Format success message
        result_parts = [
            f"**Note Created!**",
            f"",
            f"**Title:** {title}",
            f"**Section:** {section.display_name}",
        ]

        if notebook_info:
            result_parts.append(f"**Notebook:** {notebook_info.display_name}")

        result_parts.append(f"")
        result_parts.append(f"Content saved successfully.")

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        import traceback
        log(f"EXCEPTION: {e}")
        log(traceback.format_exc())
        return ChatToolResponse(error=f"Failed to create note: {str(e)}")


@app.post("/tools/search_notes", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_search_notes(request: Request):
    """
    Search for notes in OneNote.
    """
    try:
        body = await request.json()
        log(f"=== SEARCH_NOTES START ===")
        log(f"Request: {body}")

        uid = body.get("uid")
        query = body.get("query", "")
        notebook_name = body.get("notebook")
        limit = min(body.get("limit", 10), 20)

        if not uid:
            return ChatToolResponse(error="User ID is required")

        if not query:
            return ChatToolResponse(error="Search query is required")

        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        # Search pages
        pages = get_pages(uid, search_query=query, limit=limit)

        if not pages:
            return ChatToolResponse(result=f"No notes found matching '{query}'.")

        # Format results
        result_parts = [f"**Search Results for '{query}' ({len(pages)} found)**", ""]

        for i, page in enumerate(pages, 1):
            title = page.title or "Untitled"
            preview = page.preview_text or ""
            if len(preview) > 100:
                preview = preview[:100] + "..."

            # Format date
            date_str = ""
            if page.last_modified_datetime:
                try:
                    date_obj = datetime.fromisoformat(str(page.last_modified_datetime).replace('Z', '+00:00'))
                    date_str = f" - {date_obj.strftime('%b %d, %Y')}"
                except:
                    pass

            result_parts.append(f"{i}. **{title}**{date_str}")
            if preview:
                result_parts.append(f"   {preview}")
            result_parts.append("")

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        import traceback
        log(f"EXCEPTION: {e}")
        log(traceback.format_exc())
        return ChatToolResponse(error=f"Failed to search notes: {str(e)}")


@app.post("/tools/list_notebooks", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_list_notebooks(request: Request):
    """
    List the user's OneNote notebooks.
    """
    try:
        body = await request.json()
        uid = body.get("uid")

        if not uid:
            return ChatToolResponse(error="User ID is required")

        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        notebooks = get_notebooks(uid)
        if not notebooks:
            return ChatToolResponse(result="You don't have any notebooks yet. Create one in OneNote first.")

        result_parts = [f"**Your Notebooks ({len(notebooks)})**", ""]
        for i, nb in enumerate(notebooks, 1):
            default_badge = " (Default)" if nb.is_default else ""
            shared_badge = " [Shared]" if nb.is_shared else ""
            result_parts.append(f"{i}. {nb.display_name}{default_badge}{shared_badge}")

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        log(f"Error listing notebooks: {e}")
        return ChatToolResponse(error=f"Failed to list notebooks: {str(e)}")


@app.post("/tools/list_sections", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_list_sections(request: Request):
    """
    List sections in a notebook.
    """
    try:
        body = await request.json()
        uid = body.get("uid")
        notebook_name = body.get("notebook")

        if not uid:
            return ChatToolResponse(error="User ID is required")

        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        notebook_id = None
        notebook_display_name = "all notebooks"

        if notebook_name:
            notebooks = get_notebooks(uid)
            notebook_match, _ = fuzzy_match_notebook(notebook_name, notebooks)
            if notebook_match:
                notebook_id = notebook_match.id
                notebook_display_name = notebook_match.display_name
            else:
                notebook_names = [nb.display_name for nb in notebooks[:5]]
                return ChatToolResponse(
                    error=f"Could not find notebook '{notebook_name}'. Your notebooks: {', '.join(notebook_names)}"
                )

        sections = get_sections(uid, notebook_id)
        if not sections:
            return ChatToolResponse(result=f"No sections found in {notebook_display_name}.")

        result_parts = [f"**Sections in {notebook_display_name} ({len(sections)})**", ""]
        for i, s in enumerate(sections, 1):
            result_parts.append(f"{i}. {s.display_name}")

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        log(f"Error listing sections: {e}")
        return ChatToolResponse(error=f"Failed to list sections: {str(e)}")


@app.post("/tools/list_recent_notes", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_list_recent_notes(request: Request):
    """
    List recent notes/pages.
    """
    try:
        body = await request.json()
        uid = body.get("uid")
        notebook_name = body.get("notebook")
        section_name = body.get("section")
        limit = min(body.get("limit", 10), 20)

        if not uid:
            return ChatToolResponse(error="User ID is required")

        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        section_id = None
        filter_desc = ""

        if notebook_name or section_name:
            if notebook_name:
                notebooks = get_notebooks(uid)
                notebook_match, _ = fuzzy_match_notebook(notebook_name, notebooks)
                if notebook_match:
                    if section_name:
                        sections = get_sections(uid, notebook_match.id)
                        section_match, _ = fuzzy_match_section(section_name, sections)
                        if section_match:
                            section_id = section_match.id
                            filter_desc = f" in {section_match.display_name}"
                    else:
                        # Get first section of notebook
                        sections = get_sections(uid, notebook_match.id)
                        if sections:
                            section_id = sections[0].id
                            filter_desc = f" in {notebook_match.display_name}"

        pages = get_pages(uid, section_id=section_id, limit=limit)

        if not pages:
            return ChatToolResponse(result=f"No recent notes found{filter_desc}.")

        result_parts = [f"**Recent Notes{filter_desc} ({len(pages)})**", ""]
        for i, page in enumerate(pages, 1):
            title = page.title or "Untitled"
            preview = page.preview_text or ""
            if len(preview) > 80:
                preview = preview[:80] + "..."

            date_str = ""
            if page.last_modified_datetime:
                try:
                    date_obj = datetime.fromisoformat(str(page.last_modified_datetime).replace('Z', '+00:00'))
                    date_str = f" ({date_obj.strftime('%b %d')})"
                except:
                    pass

            result_parts.append(f"{i}. **{title}**{date_str}")
            if preview:
                result_parts.append(f"   {preview}")
            result_parts.append("")

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        log(f"Error listing recent notes: {e}")
        return ChatToolResponse(error=f"Failed to list recent notes: {str(e)}")


@app.post("/tools/get_note_content", tags=["chat_tools"], response_model=ChatToolResponse)
async def tool_get_note_content(request: Request):
    """
    Get the content of a specific note by searching for it.
    """
    try:
        body = await request.json()
        uid = body.get("uid")
        title = body.get("title", "")

        if not uid:
            return ChatToolResponse(error="User ID is required")

        if not title:
            return ChatToolResponse(error="Note title is required")

        access_token = get_valid_access_token(uid)
        if not access_token:
            return ChatToolResponse(error="Please connect your Microsoft account first in the app settings.")

        # Search for the note by title
        pages = get_pages(uid, search_query=title, limit=5)

        if not pages:
            return ChatToolResponse(error=f"Could not find a note with title '{title}'.")

        # Find best title match
        best_match = None
        best_score = 0.0
        title_lower = title.lower()

        for page in pages:
            page_title = (page.title or "").lower()
            score = difflib.SequenceMatcher(None, title_lower, page_title).ratio()
            if title_lower in page_title:
                score = max(score, 0.9)
            if score > best_score:
                best_score = score
                best_match = page

        if not best_match or best_score < 0.3:
            titles = [p.title for p in pages[:3]]
            return ChatToolResponse(
                error=f"Could not find note '{title}'. Did you mean: {', '.join(titles)}?"
            )

        # Get page content
        html_content = get_page_content(uid, best_match.id)
        if not html_content:
            return ChatToolResponse(error="Could not retrieve note content.")

        # Convert HTML to text
        text_content = html_to_text(html_content)

        # Truncate if too long
        if len(text_content) > 2000:
            text_content = text_content[:2000] + "\n\n... (truncated)"

        result_parts = [
            f"**{best_match.title}**",
            "",
            text_content
        ]

        return ChatToolResponse(result="\n".join(result_parts))

    except Exception as e:
        import traceback
        log(f"EXCEPTION: {e}")
        log(traceback.format_exc())
        return ChatToolResponse(error=f"Failed to get note content: {str(e)}")


# ============================================
# Omi Chat Tools Manifest
# ============================================

@app.get("/.well-known/omi-tools.json")
async def get_omi_tools_manifest():
    """
    Omi Chat Tools Manifest endpoint.

    This endpoint returns the chat tools definitions that Omi will fetch
    when the app is created or updated in the Omi App Store.
    """
    return {
        "tools": [
            {
                "name": "create_note",
                "description": "Create a new note in OneNote. Use this when the user wants to save a note, write something down, remember something, create a memo, or jot down information. The note will be created in the specified notebook/section or the default location.",
                "endpoint": "/tools/create_note",
                "method": "POST",
                "parameters": {
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "The title of the note. Defaults to 'Untitled Note' if not provided."
                        },
                        "content": {
                            "type": "string",
                            "description": "The content/body of the note. This is the main text to save. Required."
                        },
                        "notebook": {
                            "type": "string",
                            "description": "Name of the notebook to create the note in (fuzzy matched). If not provided, uses the default notebook."
                        },
                        "section": {
                            "type": "string",
                            "description": "Name of the section within the notebook (fuzzy matched). If not provided, uses the first section."
                        }
                    },
                    "required": ["content"]
                },
                "auth_required": True,
                "status_message": "Creating note in OneNote..."
            },
            {
                "name": "search_notes",
                "description": "Search for notes in OneNote. Use this when the user wants to find notes, look up something they wrote, search their notes, or find information they saved previously.",
                "endpoint": "/tools/search_notes",
                "method": "POST",
                "parameters": {
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query to find notes. Required."
                        },
                        "notebook": {
                            "type": "string",
                            "description": "Filter results to a specific notebook (fuzzy matched)."
                        },
                        "limit": {
                            "type": "integer",
                            "description": "Maximum number of results to return (default: 10, max: 20)."
                        }
                    },
                    "required": ["query"]
                },
                "auth_required": True,
                "status_message": "Searching your notes..."
            },
            {
                "name": "list_notebooks",
                "description": "List all OneNote notebooks. Use this when the user wants to see their notebooks, know what notebooks they have, or needs to choose a notebook.",
                "endpoint": "/tools/list_notebooks",
                "method": "POST",
                "parameters": {
                    "properties": {},
                    "required": []
                },
                "auth_required": True,
                "status_message": "Getting your notebooks..."
            },
            {
                "name": "list_sections",
                "description": "List sections in a notebook. Use this when the user wants to see sections in a notebook or needs to know where to save a note.",
                "endpoint": "/tools/list_sections",
                "method": "POST",
                "parameters": {
                    "properties": {
                        "notebook": {
                            "type": "string",
                            "description": "Name of the notebook to list sections from (fuzzy matched). If not provided, lists all sections."
                        }
                    },
                    "required": []
                },
                "auth_required": True,
                "status_message": "Getting sections..."
            },
            {
                "name": "list_recent_notes",
                "description": "List recent notes. Use this when the user wants to see their recent notes, latest notes, or what they've been working on.",
                "endpoint": "/tools/list_recent_notes",
                "method": "POST",
                "parameters": {
                    "properties": {
                        "notebook": {
                            "type": "string",
                            "description": "Filter to a specific notebook (fuzzy matched)."
                        },
                        "section": {
                            "type": "string",
                            "description": "Filter to a specific section (fuzzy matched). Requires notebook to be specified."
                        },
                        "limit": {
                            "type": "integer",
                            "description": "Maximum number of notes to return (default: 10, max: 20)."
                        }
                    },
                    "required": []
                },
                "auth_required": True,
                "status_message": "Getting recent notes..."
            },
            {
                "name": "get_note_content",
                "description": "Get the full content of a specific note. Use this when the user wants to read a note, see what's in a note, or retrieve the contents of a specific note.",
                "endpoint": "/tools/get_note_content",
                "method": "POST",
                "parameters": {
                    "properties": {
                        "title": {
                            "type": "string",
                            "description": "The title of the note to retrieve (fuzzy matched). Required."
                        }
                    },
                    "required": ["title"]
                },
                "auth_required": True,
                "status_message": "Getting note content..."
            }
        ]
    }


# ============================================
# Run Server
# ============================================

if __name__ == "__main__":
    import uvicorn
    port = int(os.getenv("PORT", 8080))
    uvicorn.run(app, host="0.0.0.0", port=port)