backend/routers/notebooks.py

"""
Antigravity Notebook - Notebooks Router
API endpoints for notebook management (CRUD operations).
"""

from fastapi import APIRouter, Depends, HTTPException, status
from sqlalchemy.orm import Session
from typing import List
from uuid import UUID

from backend.database import get_db, Notebook, Source, LatentTensor
from backend.models.schemas import NotebookCreate, NotebookResponse, NotebookDetail, NotebookStats
from backend.models.clara import get_clara_model
from backend.services.storage import get_storage_service
from backend.services.context_manager import get_context_manager

router = APIRouter(prefix="/notebooks", tags=["notebooks"])


@router.post("/", response_model=NotebookResponse, status_code=status.HTTP_201_CREATED)
def create_notebook(
    notebook_data: NotebookCreate,
    db: Session = Depends(get_db)
):
    """Create a new notebook"""
    notebook = Notebook(
        name=notebook_data.name,
        description=notebook_data.description
    )
    db.add(notebook)
    db.commit()
    db.refresh(notebook)

    return NotebookResponse(
        id=notebook.id,
        name=notebook.name,
        description=notebook.description,
        created_at=notebook.created_at,
        updated_at=notebook.updated_at,
        source_count=0,
        total_tokens=0
    )


@router.get("/", response_model=List[NotebookResponse])
def list_notebooks(db: Session = Depends(get_db)):
    """List all notebooks"""
    notebooks = db.query(Notebook).order_by(Notebook.created_at.desc()).all()

    response = []
    for notebook in notebooks:
        # Count sources
        source_count = db.query(Source).filter(Source.notebook_id == notebook.id).count()

        # Calculate total tokens
        sources = db.query(Source).filter(Source.notebook_id == notebook.id).all()
        total_tokens = 0
        for source in sources:
            tensors = db.query(LatentTensor).filter(LatentTensor.source_id == source.id).all()
            total_tokens += sum(t.token_count for t in tensors)

        response.append(NotebookResponse(
            id=notebook.id,
            name=notebook.name,
            description=notebook.description,
            created_at=notebook.created_at,
            updated_at=notebook.updated_at,
            source_count=source_count,
            total_tokens=total_tokens
        ))

    return response


@router.get("/{notebook_id}", response_model=NotebookDetail)
def get_notebook(notebook_id: UUID, db: Session = Depends(get_db)):
    """Get notebook details including sources"""
    notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()

    if not notebook:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Notebook {notebook_id} not found"
        )

    # Get sources with tensor counts
    sources = db.query(Source).filter(Source.notebook_id == notebook_id).all()

    from backend.models.schemas import SourceResponse

    source_responses = []
    total_tokens = 0

    for source in sources:
        tensors = db.query(LatentTensor).filter(LatentTensor.source_id == source.id).all()
        tensor_count = len(tensors)
        source_tokens = sum(t.token_count for t in tensors)
        total_tokens += source_tokens

        source_responses.append(SourceResponse(
            id=source.id,
            notebook_id=source.notebook_id,
            source_type=source.source_type,
            filename=source.filename,
            url=source.url,
            created_at=source.created_at,
            metadata=source.metadata or {},
            tensor_count=tensor_count,
            total_tokens=source_tokens
        ))

    return NotebookDetail(
        id=notebook.id,
        name=notebook.name,
        description=notebook.description,
        created_at=notebook.created_at,
        updated_at=notebook.updated_at,
        source_count=len(sources),
        total_tokens=total_tokens,
        sources=source_responses
    )


@router.get("/{notebook_id}/stats", response_model=NotebookStats)
def get_notebook_stats(notebook_id: UUID, db: Session = Depends(get_db)):
    """Get statistics about notebook context usage"""
    notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()

    if not notebook:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Notebook {notebook_id} not found"
        )

    # Get context manager
    clara = get_clara_model()
    storage = get_storage_service()
    context_mgr = get_context_manager(clara, storage)

    # Get stats
    stats = context_mgr.get_notebook_stats(db, notebook_id)

    return NotebookStats(
        notebook_id=notebook_id,
        total_sources=db.query(Source).filter(Source.notebook_id == notebook_id).count(),
        total_tensors=stats["total_segments"],
        total_tokens=stats["total_tokens"],
        max_context_tokens=stats["max_tokens"],
        context_usage_percent=stats["context_usage_percent"],
        can_fit_full_context=stats["can_fit_full_context"]
    )


@router.delete("/{notebook_id}", status_code=status.HTTP_204_NO_CONTENT)
def delete_notebook(notebook_id: UUID, db: Session = Depends(get_db)):
    """Delete a notebook and all associated data"""
    notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()

    if not notebook:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Notebook {notebook_id} not found"
        )

    # Get storage service
    storage = get_storage_service()

    # Delete all source tensors from filesystem
    sources = db.query(Source).filter(Source.notebook_id == notebook_id).all()
    for source in sources:
        storage.delete_source_tensors(db, source.id, notebook_id)

    # Delete notebook (cascades to sources, tensors, messages)
    db.delete(notebook)
    db.commit()

    return None


@router.patch("/{notebook_id}", response_model=NotebookResponse)
def update_notebook(
    notebook_id: UUID,
    notebook_data: NotebookCreate,
    db: Session = Depends(get_db)
):
    """Update notebook name/description"""
    notebook = db.query(Notebook).filter(Notebook.id == notebook_id).first()

    if not notebook:
        raise HTTPException(
            status_code=status.HTTP_404_NOT_FOUND,
            detail=f"Notebook {notebook_id} not found"
        )

    notebook.name = notebook_data.name
    if notebook_data.description is not None:
        notebook.description = notebook_data.description

    db.commit()
    db.refresh(notebook)

    # Get counts
    source_count = db.query(Source).filter(Source.notebook_id == notebook_id).count()
    sources = db.query(Source).filter(Source.notebook_id == notebook_id).all()
    total_tokens = 0
    for source in sources:
        tensors = db.query(LatentTensor).filter(LatentTensor.source_id == source.id).all()
        total_tokens += sum(t.token_count for t in tensors)

    return NotebookResponse(
        id=notebook.id,
        name=notebook.name,
        description=notebook.description,
        created_at=notebook.created_at,
        updated_at=notebook.updated_at,
        source_count=source_count,
        total_tokens=total_tokens
    )