"""
FrankenClaw — notebooklm tool

Create notebooks, add sources, generate audio overviews, infographics,
slide decks, videos, and more through Google NotebookLM.

Uses notebooklm-py (async Python client). Auth via browser cookies stored
at ~/.notebooklm/storage_state.json. Run `notebooklm login` to authenticate.

All artifact generation is async — generate returns a task_id, then poll
or wait for completion. Downloads go to /tmp/frankenclaw/ by default.
"""

import asyncio
import json
import os
from pathlib import Path

DOWNLOAD_DIR = Path("/tmp/frankenclaw")


def _ensure_download_dir():
    DOWNLOAD_DIR.mkdir(parents=True, exist_ok=True)


async def _get_client():
    """Get an authenticated NotebookLM client."""
    from notebooklm import NotebookLMClient
    try:
        return await NotebookLMClient.from_storage()
    except Exception as e:
        raise RuntimeError(
            f"NotebookLM auth failed: {e}. Cookies may have expired — "
            "run `notebooklm login` in a terminal to re-authenticate."
        )


def _ok(data):
    return json.dumps({"success": True, "result": data}, indent=2, default=str)


def _err(msg):
    return json.dumps({"success": False, "error": str(msg)}, indent=2)


async def notebooklm_manage(
    action: str,
    notebook_id: str = "",
    title: str = "",
) -> str:
    """
    Manage NotebookLM notebooks: list, create, get, rename, delete.

    Args:
        action: One of: list, create, get, rename, delete.
        notebook_id: Notebook ID (required for get, rename, delete).
        title: Notebook title (required for create, rename).

    Returns:
        JSON with notebook info.
    """
    try:
        async with await _get_client() as client:
            if action == "list":
                notebooks = await client.notebooks.list()
                return _ok([{"id": nb.id, "title": nb.title} for nb in notebooks])

            elif action == "create":
                if not title:
                    return _err("'create' requires a title.")
                nb = await client.notebooks.create(title)
                return _ok({"id": nb.id, "title": nb.title})

            elif action == "get":
                if not notebook_id:
                    return _err("'get' requires a notebook_id.")
                nb = await client.notebooks.get(notebook_id)
                return _ok({"id": nb.id, "title": nb.title})

            elif action == "rename":
                if not notebook_id or not title:
                    return _err("'rename' requires notebook_id and title.")
                await client.notebooks.rename(notebook_id, title)
                return _ok({"renamed": notebook_id, "new_title": title})

            elif action == "delete":
                if not notebook_id:
                    return _err("'delete' requires a notebook_id.")
                await client.notebooks.delete(notebook_id)
                return _ok({"deleted": notebook_id})

            else:
                return _err(f"Unknown action '{action}'. Use: list, create, get, rename, delete.")
    except Exception as e:
        return _err(e)


async def notebooklm_source(
    action: str,
    notebook_id: str,
    url: str = "",
    text: str = "",
    file_path: str = "",
    source_id: str = "",
) -> str:
    """
    Manage sources in a NotebookLM notebook: add URLs, text, or files; list or delete sources.

    Args:
        action: One of: add_url, add_text, add_file, list, delete, get_content.
        notebook_id: Target notebook ID.
        url: URL to add (for add_url).
        text: Text content to add (for add_text).
        file_path: Local file path to upload (for add_file — PDF, TXT, etc.).
        source_id: Source ID (for delete, get_content).

    Returns:
        JSON with source info.
    """
    if not notebook_id:
        return _err("notebook_id is required.")

    try:
        async with await _get_client() as client:
            if action == "add_url":
                if not url:
                    return _err("'add_url' requires a url.")
                src = await client.sources.add_url(notebook_id, url, wait=True)
                return _ok({"id": src.id, "title": getattr(src, "title", url)})

            elif action == "add_text":
                if not text:
                    return _err("'add_text' requires text content.")
                src = await client.sources.add_text(notebook_id, text)
                return _ok({"id": src.id, "title": getattr(src, "title", "text")})

            elif action == "add_file":
                if not file_path:
                    return _err("'add_file' requires a file_path.")
                src = await client.sources.add_file(notebook_id, file_path)
                return _ok({"id": src.id, "title": getattr(src, "title", file_path)})

            elif action == "list":
                sources = await client.sources.list(notebook_id)
                return _ok([{"id": s.id, "title": getattr(s, "title", "?")} for s in sources])

            elif action == "delete":
                if not source_id:
                    return _err("'delete' requires a source_id.")
                await client.sources.delete(notebook_id, source_id)
                return _ok({"deleted": source_id})

            elif action == "get_content":
                if not source_id:
                    return _err("'get_content' requires a source_id.")
                text = await client.sources.get_fulltext(notebook_id, source_id)
                return _ok({"source_id": source_id, "content": str(text)[:5000]})

            else:
                return _err(f"Unknown action '{action}'. Use: add_url, add_text, add_file, list, delete, get_content.")
    except Exception as e:
        return _err(e)


async def notebooklm_generate(
    artifact_type: str,
    notebook_id: str,
    instructions: str = "",
    wait: bool = True,
) -> str:
    """
    Generate content from a NotebookLM notebook: audio podcasts, videos,
    infographics, slide decks, reports, quizzes, flashcards, mind maps, data tables.

    Args:
        artifact_type: One of: audio, video, cinematic_video, infographic, slide_deck,
                       report, study_guide, quiz, flashcards, mind_map, data_table.
        notebook_id: Source notebook ID.
        instructions: Optional natural language instructions (e.g. "make it fun and engaging",
                      "focus on the marketing strategy", "keep it under 5 minutes").
        wait: If true (default), waits for generation to complete before returning.

    Returns:
        JSON with generation status and artifact info.
    """
    if not notebook_id:
        return _err("notebook_id is required.")

    generators = {
        "audio": "generate_audio",
        "video": "generate_video",
        "cinematic_video": "generate_cinematic_video",
        "infographic": "generate_infographic",
        "slide_deck": "generate_slide_deck",
        "report": "generate_report",
        "study_guide": "generate_study_guide",
        "quiz": "generate_quiz",
        "flashcards": "generate_flashcards",
        "mind_map": "generate_mind_map",
        "data_table": "generate_data_table",
    }

    if artifact_type not in generators:
        return _err(f"Unknown type '{artifact_type}'. Use: {', '.join(generators.keys())}")

    try:
        async with await _get_client() as client:
            method = getattr(client.artifacts, generators[artifact_type])

            # Pass instructions if the method accepts them
            kwargs = {}
            if instructions:
                import inspect
                sig = inspect.signature(method)
                if "instructions" in sig.parameters:
                    kwargs["instructions"] = instructions

            status = await method(notebook_id, **kwargs)
            task_id = getattr(status, "task_id", None) or getattr(status, "id", None)

            if wait and task_id:
                try:
                    result = await client.artifacts.wait_for_completion(
                        notebook_id, task_id, timeout=300
                    )
                    return _ok({
                        "artifact_type": artifact_type,
                        "task_id": str(task_id),
                        "status": "completed",
                        "result": str(result),
                    })
                except TimeoutError:
                    return _ok({
                        "artifact_type": artifact_type,
                        "task_id": str(task_id),
                        "status": "still_generating",
                        "message": "Generation in progress. Use notebooklm_download to check later.",
                    })

            return _ok({
                "artifact_type": artifact_type,
                "task_id": str(task_id),
                "status": "started",
            })
    except Exception as e:
        return _err(e)


async def notebooklm_download(
    artifact_type: str,
    notebook_id: str,
    output_path: str = "",
) -> str:
    """
    Download a generated artifact from a NotebookLM notebook.

    Args:
        artifact_type: One of: audio, video, infographic, slide_deck, report,
                       mind_map, data_table, quiz, flashcards.
        notebook_id: Notebook ID containing the artifact.
        output_path: Where to save (default: /tmp/frankenclaw/<type>.<ext>).

    Returns:
        JSON with the file path of the downloaded artifact.
    """
    if not notebook_id:
        return _err("notebook_id is required.")

    downloaders = {
        "audio": ("download_audio", "mp3"),
        "video": ("download_video", "mp4"),
        "infographic": ("download_infographic", "png"),
        "slide_deck": ("download_slide_deck", "pdf"),
        "report": ("download_report", "md"),
        "mind_map": ("download_mind_map", "json"),
        "data_table": ("download_data_table", "csv"),
        "quiz": ("download_quiz", "md"),
        "flashcards": ("download_flashcards", "md"),
    }

    if artifact_type not in downloaders:
        return _err(f"Unknown type '{artifact_type}'. Use: {', '.join(downloaders.keys())}")

    try:
        _ensure_download_dir()
        method_name, ext = downloaders[artifact_type]

        if not output_path:
            output_path = str(DOWNLOAD_DIR / f"{artifact_type}.{ext}")

        async with await _get_client() as client:
            method = getattr(client.artifacts, method_name)
            await method(notebook_id, output_path)
            size = os.path.getsize(output_path)
            return _ok({
                "artifact_type": artifact_type,
                "path": output_path,
                "size_bytes": size,
            })
    except Exception as e:
        return _err(e)


async def notebooklm_ask(
    notebook_id: str,
    question: str,
) -> str:
    """
    Ask a question about the sources in a NotebookLM notebook. Uses the
    notebook's AI to generate an answer grounded in the loaded sources.

    Args:
        notebook_id: Notebook ID to query.
        question: The question to ask.

    Returns:
        JSON with the AI's answer.
    """
    if not notebook_id or not question:
        return _err("Both notebook_id and question are required.")

    try:
        async with await _get_client() as client:
            result = await client.chat.ask(notebook_id, question)
            return _ok({
                "answer": getattr(result, "answer", str(result)),
                "notebook_id": notebook_id,
            })
    except Exception as e:
        return _err(e)
