"""Async NotebookLM work: notebook + sources + wait + chat (via notebooklm-py)."""

from __future__ import annotations

import asyncio
import logging
from datetime import datetime, timezone
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from notebooklm import NotebookLMClient
    from folio_notion.pipeline_context import PipelineContext

log = logging.getLogger(__name__)


async def _open_client(ctx: PipelineContext) -> NotebookLMClient:
    try:
        from notebooklm import NotebookLMClient as NLC
    except ImportError as e:
        raise RuntimeError(
            'NotebookLM integration needs: pip install "folio-notion[notebooklm]"'
        ) from e

    storage = ctx.notebooklm_storage
    if not storage or not storage.is_file():
        raise RuntimeError(
            "Missing NotebookLM session. Run: fn connect notebooklm"
        )

    path_arg = str(storage.resolve())
    log.info("NotebookLM: opening client (storage=%s)", path_arg)
    return await NLC.from_storage(path_arg)


async def _resolve_notebook(client: NotebookLMClient, ctx: PipelineContext) -> str:
    nid = ctx.notebooklm_notebook_id
    if nid:
        log.info("NotebookLM: using existing notebook_id=%s", nid)
        await client.notebooks.get(nid)
        return nid

    title = ctx.run_title or f"Folio {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M')} UTC"
    log.info("NotebookLM: creating notebook %r", title)
    nb = await client.notebooks.create(title)
    log.info("NotebookLM: created notebook_id=%s", nb.id)
    return nb.id


async def _upload_pdfs(
    client: NotebookLMClient, ctx: PipelineContext, notebook_id: str
) -> list[str]:
    source_ids: list[str] = []
    timeout = ctx.notebooklm_source_upload_timeout_sec
    for i, pdf in enumerate(ctx.pdf_paths, 1):
        log.info(
            "NotebookLM: upload %d/%d %s (timeout=%ss)",
            i,
            len(ctx.pdf_paths),
            pdf.name,
            int(timeout),
        )
        src = await client.sources.add_file(
            notebook_id,
            pdf,
            wait=True,
            wait_timeout=float(timeout),
        )
        source_ids.append(src.id)
    return source_ids


async def _wait_sources_ready(
    client: NotebookLMClient, ctx: PipelineContext, notebook_id: str, source_ids: list[str]
) -> None:
    if not source_ids:
        return
    timeout = float(ctx.notebooklm_source_upload_timeout_sec)
    log.info(
        "NotebookLM: wait_for_sources (%d sources, timeout=%ss)",
        len(source_ids),
        int(timeout),
    )
    await client.sources.wait_for_sources(
        notebook_id,
        source_ids,
        timeout=timeout,
    )


async def _ask_synthesis(
    client: NotebookLMClient, ctx: PipelineContext, notebook_id: str
) -> str:
    retries = max(0, int(ctx.notebooklm_chat_retries))
    delay = float(ctx.notebooklm_chat_retry_delay_sec)
    last: BaseException | None = None
    for attempt in range(retries + 1):
        log.info(
            "NotebookLM: chat.ask (attempt %d/%d, prompt_len=%d)",
            attempt + 1,
            retries + 1,
            len(ctx.synthesis_prompt),
        )
        try:
            result = await client.chat.ask(notebook_id, ctx.synthesis_prompt)
            return result.answer
        except Exception as e:
            last = e
            log.warning("NotebookLM: chat.ask failed: %s", e)
            if attempt >= retries:
                raise
            await asyncio.sleep(delay * (attempt + 1))
    assert last is not None
    raise last


async def run_notebooklm(ctx: PipelineContext) -> None:
    client = await _open_client(ctx)
    async with client:
        notebook_id = await _resolve_notebook(client, ctx)
        ctx.notebooklm_notebook_id = notebook_id

        source_ids = await _upload_pdfs(client, ctx, notebook_id)

        settle = float(ctx.notebooklm_post_upload_wait_sec)
        if settle > 0:
            log.info("NotebookLM: post-upload settle sleep %ss", settle)
            await asyncio.sleep(settle)

        await _wait_sources_ready(client, ctx, notebook_id, source_ids)

        ctx.synthesis_text = await _ask_synthesis(client, ctx, notebook_id)
        log.info(
            "NotebookLM: synthesis received (%d chars)",
            len(ctx.synthesis_text or ""),
        )
