"""
MCP Resources for reMarkable tablet access.

Provides:
- remarkable:///{path}.txt - extracted text from any document
- remarkableraw:///{path}.txt - raw PDF/EPUB text content (SSH mode only, enumerated)

Resources are loaded at startup (SSH) or in background batches (cloud).
Respects REMARKABLE_ROOT_PATH environment variable for folder filtering.
"""

import asyncio
import logging
import os
import tempfile
from pathlib import Path
from typing import Optional, Set

from remarkable_mcp.server import mcp

logger = logging.getLogger(__name__)


def _get_root_path() -> str:
    """Get the configured root path filter, or '/' for full access.

    Handles: empty string, '/', '/Work', '/Work/', 'Work' -> normalized path
    """
    root = os.environ.get("REMARKABLE_ROOT_PATH", "").strip()
    # Empty or "/" means full access
    if not root or root == "/":
        return "/"
    # Normalize: ensure starts with / and no trailing slash
    if not root.startswith("/"):
        root = "/" + root
    if root.endswith("/"):
        root = root.rstrip("/")
    return root


def _is_within_root(path: str, root: str) -> bool:
    """Check if a path is within the configured root (case-insensitive)."""
    if root == "/":
        return True
    # Path must equal root or be a child of root (case-insensitive)
    path_lower = path.lower()
    root_lower = root.lower()
    return path_lower == root_lower or path_lower.startswith(root_lower + "/")


def _apply_root_filter(path: str, root: str) -> str:
    """Apply root filter to a path for display/URI purposes.

    If root is '/Work', then '/Work/Project' becomes '/Project' in output.
    Case-insensitive matching, preserves original case in output.
    """
    if root == "/":
        return path
    path_lower = path.lower()
    root_lower = root.lower()
    if path_lower == root_lower:
        return "/"
    if path_lower.startswith(root_lower + "/"):
        return path[len(root) :]
    return path


# Background loader state
_registered_docs: Set[str] = set()  # Track document IDs for text resources
_registered_raw: Set[str] = set()  # Track document IDs for raw resources
_registered_uris: Set[str] = set()  # Track URIs for collision detection


def _is_ssh_mode() -> bool:
    """Check if SSH transport is enabled (evaluated at runtime)."""
    return os.environ.get("REMARKABLE_USE_SSH", "").lower() in ("1", "true", "yes")


def _make_doc_resource(client, document):
    """Create a resource function for a document.

    Returns only user-supplied content: typed text, annotations, highlights,
    and OCR for handwritten content. Does NOT include original PDF/EPUB text.
    Use raw resources for original document text.
    """
    from remarkable_mcp.extract import extract_text_from_document_zip

    def doc_resource() -> str:
        try:
            text_parts = []

            # Download notebook data for annotations/typed text/handwritten
            raw = client.download(document)
            with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as tmp:
                tmp.write(raw)
                tmp_path = Path(tmp.name)
            try:
                # First try without OCR (faster)
                content = extract_text_from_document_zip(tmp_path, include_ocr=False)

                if content["typed_text"]:
                    text_parts.extend(content["typed_text"])
                if content["highlights"]:
                    if text_parts:
                        text_parts.append("\n--- Highlights ---")
                    text_parts.extend(content["highlights"])

                # If no text found and document has pages, try OCR for handwritten
                if not text_parts and content["pages"] > 0:
                    content = extract_text_from_document_zip(tmp_path, include_ocr=True)
                    if content["handwritten_text"]:
                        text_parts.extend(content["handwritten_text"])

                return "\n\n".join(text_parts) if text_parts else "(No user content)"
            finally:
                tmp_path.unlink(missing_ok=True)
        except Exception as e:
            return f"Error: {e}"

    return doc_resource


def _make_raw_resource(client, document, file_type: str):
    """Create a resource function for raw PDF/EPUB text extraction."""
    from remarkable_mcp.api import download_raw_file
    from remarkable_mcp.extract import extract_text_from_epub, extract_text_from_pdf

    def raw_resource() -> str:
        try:
            if not _is_ssh_mode():
                return "Error: Raw file download only available in SSH mode"

            raw_data = download_raw_file(client, document, file_type)

            if not raw_data:
                return f"Raw {file_type.upper()} file not found"

            # Extract text from the raw file
            with tempfile.NamedTemporaryFile(suffix=f".{file_type}", delete=False) as tmp:
                tmp.write(raw_data)
                tmp_path = Path(tmp.name)

            try:
                if file_type == "pdf":
                    text = extract_text_from_pdf(tmp_path)
                elif file_type == "epub":
                    text = extract_text_from_epub(tmp_path)
                else:
                    text = f"Unsupported file type: {file_type}"

                return text if text else f"(No text content in {file_type.upper()} file)"
            finally:
                tmp_path.unlink(missing_ok=True)
        except Exception as e:
            return f"Error: {e}"

    return raw_resource


def _register_document(
    client, doc, items_by_id=None, file_types: dict = None, root: str = "/"
) -> bool:
    """Register a single document as a text resource (and raw resource if PDF/EPUB).

    Args:
        client: The reMarkable API client
        doc: Document metadata object
        items_by_id: Dict mapping IDs to items for path resolution
        file_types: Dict mapping doc IDs to file types (for raw resources)
        root: Root path filter (documents outside root are skipped)
    """
    global _registered_docs, _registered_raw, _registered_uris

    doc_id = doc.ID

    # Skip if already registered (by ID)
    if doc_id in _registered_docs:
        return False

    # Skip cloud-archived documents (not available on device)
    if hasattr(doc, "is_cloud_archived") and doc.is_cloud_archived:
        return False

    # Get the full path
    doc_name = doc.VissibleName
    if items_by_id:
        from remarkable_mcp.api import get_item_path

        full_path = get_item_path(doc, items_by_id)
    else:
        full_path = f"/{doc_name}"

    # Filter by root path
    if not _is_within_root(full_path, root):
        return False

    # Apply root filter for display paths (e.g., /Work/Project -> /Project)
    display_path = _apply_root_filter(full_path, root)

    # Use the filtered path for URIs
    uri_path = display_path.lstrip("/")

    # Register text resource (use /// for empty netloc)
    base_uri = f"remarkable:///{uri_path}.txt"
    counter = 1
    final_uri = base_uri
    display_name = f"{display_path}.txt"
    while final_uri in _registered_uris:
        final_uri = f"remarkable:///{uri_path}_{counter}.txt"
        display_name = f"{display_path} ({counter}).txt"
        counter += 1

    desc = f"Content from '{display_path}'"
    if doc.ModifiedClient:
        desc += f" (modified: {doc.ModifiedClient})"

    mcp.resource(final_uri, name=display_name, description=desc, mime_type="text/plain")(
        _make_doc_resource(client, doc)
    )

    _registered_docs.add(doc_id)
    _registered_uris.add(final_uri)

    # Also register raw resource for PDF/EPUB files (SSH mode only)
    if _is_ssh_mode() and file_types is not None:
        # Use pre-loaded file types (fast)
        file_type = file_types.get(doc_id)
        if file_type in ("pdf", "epub"):
            # Raw resources now return extracted text, use .txt extension
            raw_uri = f"remarkableraw:///{uri_path}.{file_type}.txt"
            raw_counter = 1
            final_raw_uri = raw_uri
            raw_display = f"{display_path} (raw {file_type.upper()}).txt"
            while final_raw_uri in _registered_uris:
                final_raw_uri = f"remarkableraw:///{uri_path}_{raw_counter}.{file_type}.txt"
                raw_display = f"{display_path} (raw {file_type.upper()}) ({raw_counter}).txt"
                raw_counter += 1

            raw_desc = f"Raw {file_type.upper()} text content: '{display_path}'"
            if doc.ModifiedClient:
                raw_desc += f" (modified: {doc.ModifiedClient})"

            mcp.resource(
                final_raw_uri,
                name=raw_display,
                description=raw_desc,
                mime_type="text/plain",
            )(_make_raw_resource(client, doc, file_type))

            _registered_raw.add(doc_id)
            _registered_uris.add(final_raw_uri)

    return True


def load_all_documents_sync() -> int:
    """
    Load and register all documents synchronously.
    Used for SSH mode where loading is fast.
    Returns the number of documents registered.

    Respects REMARKABLE_ROOT_PATH environment variable.
    """
    global _registered_docs, _registered_raw

    from remarkable_mcp.api import get_items_by_id, get_rmapi

    client = get_rmapi()
    items = client.get_meta_items()
    items_by_id = get_items_by_id(items)
    documents = [item for item in items if not item.is_folder]

    root = _get_root_path()
    if root != "/":
        logger.info(f"Root path filter: {root}")

    logger.info(f"Found {len(documents)} documents")

    # Pre-load all file types in a single SSH call (SSH mode optimization)
    file_types = {}
    if _is_ssh_mode() and hasattr(client, "get_all_file_types"):
        logger.info("Pre-loading file types for raw resources...")
        file_types = client.get_all_file_types()
        logger.info(f"Loaded {len(file_types)} file types")

    for doc in documents:
        try:
            _register_document(
                client,
                doc,
                items_by_id,
                file_types if _is_ssh_mode() else None,
                root=root,
            )
        except Exception as e:
            logger.debug(f"Failed to register '{doc.VissibleName}': {e}")

    logger.info(
        f"Registered {len(_registered_docs)} text resources"
        + (f", {len(_registered_raw)} raw resources (PDF/EPUB)" if _registered_raw else "")
        + (f" (filtered to {root})" if root != "/" else "")
    )
    return len(_registered_docs)


async def _load_documents_background(shutdown_event: asyncio.Event):
    """
    Background task to load and register documents in batches.
    Used for Cloud mode only - SSH mode uses load_all_documents_sync().

    Respects REMARKABLE_ROOT_PATH environment variable.
    """
    try:
        from remarkable_mcp.api import get_items_by_id, get_rmapi

        client = get_rmapi()
        loop = asyncio.get_event_loop()

        batch_size = 10
        offset = 0
        consecutive_errors = 0
        max_consecutive_errors = 3
        items_by_id = {}  # Build incrementally

        root = _get_root_path()
        if root != "/":
            logger.info(f"Root path filter: {root}")

        while True:
            # Check for shutdown
            if shutdown_event.is_set():
                logger.info("Background document loader cancelled by shutdown")
                break

            # Fetch next batch - run sync code in executor to not block
            try:
                items = await loop.run_in_executor(
                    None, lambda: client.get_meta_items(limit=offset + batch_size)
                )
                # Update items_by_id with all items for path resolution
                items_by_id = get_items_by_id(items)
                consecutive_errors = 0  # Reset on success
            except Exception as e:
                consecutive_errors += 1
                logger.warning(f"Error fetching documents (attempt {consecutive_errors}): {e}")
                if consecutive_errors >= max_consecutive_errors:
                    logger.error(
                        f"Background loader stopping after {max_consecutive_errors} "
                        "consecutive errors"
                    )
                    break
                # Wait before retry
                await asyncio.sleep(2**consecutive_errors)
                continue

            # Get documents from this batch (skip folders and already-fetched)
            documents = [item for item in items if not item.is_folder]
            batch_docs = documents[offset : offset + batch_size]

            if not batch_docs:
                # No more documents
                logger.info(
                    f"Background loader complete: {len(_registered_docs)} documents registered"
                    + (f" (filtered to {root})" if root != "/" else "")
                )
                break

            # Register this batch (no file_types in cloud mode - raw resources not available)
            registered_count = 0
            for doc in batch_docs:
                if shutdown_event.is_set():
                    break
                try:
                    if _register_document(client, doc, items_by_id, file_types=None, root=root):
                        registered_count += 1
                except Exception as e:
                    logger.debug(f"Failed to register document '{doc.VissibleName}': {e}")

            if registered_count > 0:
                logger.debug(
                    f"Registered batch of {registered_count} documents "
                    f"(total: {len(_registered_docs)})"
                )

            offset += batch_size

            # Yield control - allow other async tasks to run
            # Small delay to be gentle on the API
            await asyncio.sleep(0.1)

    except asyncio.CancelledError:
        logger.info("Background document loader cancelled")
        raise
    except Exception as e:
        logger.warning(f"Background document loader error: {e}")


def start_background_loader() -> Optional[asyncio.Task]:
    """Start the background document loader task. Returns the task."""
    shutdown_event = asyncio.Event()

    try:
        task = asyncio.create_task(_load_documents_background(shutdown_event))
        # Store the event on the task so we can access it later
        task.shutdown_event = shutdown_event  # type: ignore[attr-defined]
        logger.info("Started background document loader")
        return task
    except Exception as e:
        logger.warning(f"Could not start background loader: {e}")
        return None


async def stop_background_loader(task: Optional[asyncio.Task]):
    """Stop the background document loader task."""
    if task is None:
        return

    # Signal shutdown via event
    if hasattr(task, "shutdown_event"):
        task.shutdown_event.set()

    # Cancel and wait
    if not task.done():
        task.cancel()
        try:
            await task
        except asyncio.CancelledError:
            pass

    logger.info("Stopped background document loader")
