"""Jupyter notebook execution endpoints.

Provides per-cell execution with multi-session support. Each session gets its
own Jupyter kernel via nbclient. Requires the ``notebooks`` optional extra::

    pip install open-terminal[notebooks]
"""

import asyncio
import json
import os
import time
import uuid
from typing import Optional

import aiofiles
import aiofiles.os
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, Field

import nbformat
from nbclient import NotebookClient


# ---------------------------------------------------------------------------
# Session manager
# ---------------------------------------------------------------------------

_IDLE_TIMEOUT = 30 * 60  # 30 minutes


class _Session:
    """Wraps a NotebookClient for a specific notebook."""

    __slots__ = ("id", "path", "nb", "client", "busy", "created_at", "last_used")

    def __init__(self, session_id: str, path: str, nb, client):
        self.id = session_id
        self.path = path
        self.nb = nb
        self.client = client
        self.busy = False
        self.created_at = time.time()
        self.last_used = time.time()


_sessions: dict[str, _Session] = {}
_cleanup_task: Optional[asyncio.Task] = None


async def _idle_cleanup_loop():
    """Periodically remove sessions idle for more than _IDLE_TIMEOUT."""
    while True:
        await asyncio.sleep(60)
        now = time.time()
        stale = [
            sid
            for sid, s in _sessions.items()
            if now - s.last_used > _IDLE_TIMEOUT and not s.busy
        ]
        for sid in stale:
            await _destroy_session(sid)


async def _destroy_session(session_id: str):
    session = _sessions.pop(session_id, None)
    if session and session.client:
        try:
            await session.client._async_cleanup_kernel()
        except Exception:
            pass


def _ensure_cleanup_task():
    global _cleanup_task
    if _cleanup_task is None or _cleanup_task.done():
        _cleanup_task = asyncio.create_task(_idle_cleanup_loop())


# ---------------------------------------------------------------------------
# Request / response models
# ---------------------------------------------------------------------------


class CreateSessionRequest(BaseModel):
    path: str = Field(..., description="Absolute path to the .ipynb file.")


class CreateSessionResponse(BaseModel):
    id: str
    kernel: str
    status: str


class ExecuteCellRequest(BaseModel):
    cell_index: int = Field(..., description="Zero-based cell index to execute.")
    source: Optional[str] = Field(
        None, description="Override cell source. If omitted, uses the source already in the notebook."
    )


class ExecuteCellResponse(BaseModel):
    status: str
    execution_count: Optional[int] = None
    outputs: list = Field(default_factory=list)


class SessionStatusResponse(BaseModel):
    id: str
    path: str
    kernel: str
    status: str


# ---------------------------------------------------------------------------
# Router
# ---------------------------------------------------------------------------


def create_notebooks_router(verify_api_key) -> APIRouter:
    """Create the notebooks router with the given auth dependency."""

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



    @router.post(
        "",
        response_model=CreateSessionResponse,
        operation_id="create_notebook_session",
        summary="Create a notebook session",
        description="Start a Jupyter kernel for the given notebook. Returns a session ID for subsequent execute calls.",
        include_in_schema=False,
    )
    async def create_session(req: CreateSessionRequest):

        _ensure_cleanup_task()

        path = os.path.abspath(req.path)
        if not await aiofiles.os.path.isfile(path):
            raise HTTPException(status_code=404, detail=f"Notebook not found: {path}")

        # Read and parse the notebook
        async with aiofiles.open(path, encoding="utf-8") as f:
            content = await f.read()

        try:
            nb = nbformat.reads(content, as_version=4)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid notebook: {e}")

        kernel_name = nb.metadata.get("kernelspec", {}).get("name", "python3")

        # Start kernel — follow nbclient's setup_kernel pattern
        client = NotebookClient(nb, kernel_name=kernel_name, timeout=120)
        try:
            client.create_kernel_manager()
            await client.async_start_new_kernel()
            await client.async_start_new_kernel_client()
        except Exception as e:
            raise HTTPException(
                status_code=500,
                detail=f"Failed to start kernel '{kernel_name}': {e}",
            )

        session_id = uuid.uuid4().hex[:12]
        _sessions[session_id] = _Session(session_id, path, nb, client)

        return CreateSessionResponse(
            id=session_id, kernel=kernel_name, status="ready"
        )

    @router.post(
        "/{session_id}/execute",
        response_model=ExecuteCellResponse,
        operation_id="execute_notebook_cell",
        summary="Execute a notebook cell",
        description="Execute a single cell in the given session. Optionally override the cell source. "
        "Updates the .ipynb file in place after execution.",
        include_in_schema=False,
    )
    async def execute_cell(session_id: str, req: ExecuteCellRequest):


        session = _sessions.get(session_id)
        if not session:
            raise HTTPException(status_code=404, detail="Session not found")
        if session.busy:
            raise HTTPException(status_code=409, detail="Cell already executing")

        nb = session.nb
        if req.cell_index < 0 or req.cell_index >= len(nb.cells):
            raise HTTPException(
                status_code=400,
                detail=f"cell_index {req.cell_index} out of range (0..{len(nb.cells) - 1})",
            )

        cell = nb.cells[req.cell_index]
        if req.source is not None:
            cell.source = req.source

        session.busy = True
        session.last_used = time.time()

        try:
            await session.client.async_execute_cell(cell, req.cell_index)
        except Exception as e:
            session.busy = False
            # Return the error as a cell output rather than HTTP error
            return ExecuteCellResponse(
                status="error",
                outputs=[
                    {
                        "output_type": "error",
                        "ename": type(e).__name__,
                        "evalue": str(e),
                        "traceback": [str(e)],
                    }
                ],
            )

        session.busy = False
        session.last_used = time.time()

        # Serialize outputs
        outputs = []
        for o in cell.outputs:
            od = dict(o)
            if "data" in od:
                od["data"] = dict(od["data"])
            outputs.append(od)

        ec = cell.get("execution_count")

        # Save notebook to disk
        try:
            nb_json = nbformat.writes(nb)
            async with aiofiles.open(session.path, "w", encoding="utf-8") as f:
                await f.write(nb_json)
        except Exception:
            pass  # non-fatal

        return ExecuteCellResponse(
            status="ok", execution_count=ec, outputs=outputs
        )

    @router.get(
        "/{session_id}",
        response_model=SessionStatusResponse,
        operation_id="get_notebook_session",
        summary="Get notebook session status",
        include_in_schema=False,
    )
    async def get_session(session_id: str):
        session = _sessions.get(session_id)
        if not session:
            raise HTTPException(status_code=404, detail="Session not found")

        kernel_name = session.nb.metadata.get("kernelspec", {}).get(
            "name", "python3"
        )
        status = "busy" if session.busy else "ready"
        return SessionStatusResponse(
            id=session.id, path=session.path, kernel=kernel_name, status=status
        )

    @router.delete(
        "/{session_id}",
        operation_id="delete_notebook_session",
        summary="Stop a notebook session",
        include_in_schema=False,
    )
    async def delete_session(session_id: str):
        if session_id not in _sessions:
            raise HTTPException(status_code=404, detail="Session not found")
        await _destroy_session(session_id)
        return {"status": "stopped"}

    return router
