"""Knowledge Base documents resources — upload, ingest, polling."""

from __future__ import annotations

import asyncio
import io
import logging
import time
import zipfile
from pathlib import Path

from ..._api.client import AsyncBaseResource, BaseResource
from ..._api.generated.knowbase.knowbase__api__knowbase.client import (
    KnowbaseKnowbaseAPI as _AsyncAPI,
)
from ..._api.generated.knowbase.knowbase__api__knowbase.models import (
    ArchiveBulkCreateResponse,
    ArchiveJobStatus,
    Document,
    DocumentUploadRequest,
    KnowbaseProjectsDocumentsBulkCreateRequest,
    PaginatedDocumentList,
)
from ..._api.generated.knowbase.knowbase__api__knowbase.sync_client import (
    SyncKnowbaseKnowbaseAPI as _SyncAPI,
)
from ..._config import SDKConfig
from ...exceptions import api_error_handler, async_api_error_handler
from .._polling import AsyncPollingMixin, PollingMixin

logger = logging.getLogger(__name__)


class _DocumentsResource(PollingMixin, BaseResource):
    """Documents within a KB project (sync)."""

    def __init__(self, config: SDKConfig, project_slug: str) -> None:
        super().__init__(config, use_upload_timeout=True)
        self._api = _SyncAPI(self._http_client)
        self._slug = project_slug

    @api_error_handler
    def list(self, *, page: int = 1, page_size: int = 20) -> PaginatedDocumentList:
        """List documents in this project."""
        return self._api.projects_documents_list(  # type: ignore[return-value]
            project_slug=self._slug, page=page, page_size=page_size,
        )

    @api_error_handler
    def upload(
        self,
        *,
        title: str,
        content: str,
        source_url: str | None = None,
    ) -> Document:
        """Upload a document (Markdown/plain text).

        The document will be chunked and embedded asynchronously in the background.

        Args:
            title: Document title.
            content: Raw Markdown or plain text content.
            source_url: Optional original URL (e.g. link to the source page).
        """
        logger.debug("KB document upload: project=%s title=%r", self._slug, title)
        data = DocumentUploadRequest(title=title, content=content, source_url=source_url)
        result = self._api.projects_documents_create(project_slug=self._slug, data=data)
        logger.info("KB document uploaded: project=%s id=%s status=%s", self._slug, result.id, result.processing_status)
        return result

    @api_error_handler
    def upload_file(self, path: Path, *, source_url: str | None = None) -> Document:
        """Upload a local Markdown or text file.

        Args:
            path: Path to a .md or .txt file.
            source_url: Optional URL to associate with this document.
        """
        path = Path(path)
        content = path.read_text(encoding="utf-8")
        title = path.stem.replace("-", " ").replace("_", " ").title()
        logger.debug("KB document upload_file: project=%s path=%s", self._slug, path)
        return self.upload(title=title, content=content, source_url=source_url)

    @api_error_handler
    def get(self, document_id: str) -> Document:
        """Get a document by ID."""
        return self._api.projects_documents_retrieve(id=document_id, project_slug=self._slug)

    @api_error_handler
    def delete(self, document_id: str) -> None:
        """Delete a document and all its chunks."""
        logger.info("KB document delete: project=%s id=%s", self._slug, document_id)
        self._api.projects_documents_destroy(id=document_id, project_slug=self._slug)

    @api_error_handler
    def upload_archive(self, zip_path: str | Path) -> ArchiveBulkCreateResponse:
        """Upload a ZIP archive for bulk document ingestion.

        Each file inside the ZIP is converted to Markdown, LLM-summarized,
        deduplicated, and indexed asynchronously by the background worker.

        Args:
            zip_path: Path to a ``.zip`` file on disk.

        Returns:
            :class:`ArchiveBulkCreateResponse` with ``job_id`` for polling.
            Use :meth:`wait_for_archive_done` to block until ingestion completes.
        """
        zip_path = Path(zip_path)
        logger.debug("KB upload_archive: project=%s path=%s", self._slug, zip_path)
        with zip_path.open("rb") as fh:
            data = KnowbaseProjectsDocumentsBulkCreateRequest(archive=(zip_path.name, fh, "application/zip"))
            result = self._api.projects_documents_bulk_create(project_slug=self._slug, data=data)
        logger.info("KB upload_archive queued: project=%s job_id=%s", self._slug, result.job_id)
        return result

    @api_error_handler
    def upload_dir(self, dir_path: str | Path, *, prefix: str = "") -> ArchiveBulkCreateResponse:
        """Zip a local directory in-memory and upload it for bulk document ingestion.

        Args:
            dir_path: Path to a local directory to zip and upload.
            prefix: Optional path prefix inside the ZIP (default: none).

        Returns:
            :class:`ArchiveBulkCreateResponse` with ``job_id`` for polling.
        """
        dir_path = Path(dir_path).resolve()
        if not dir_path.is_dir():
            raise ValueError(f"Not a directory: {dir_path}")

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for file in sorted(dir_path.rglob("*")):
                if file.is_file():
                    arcname = str(file.relative_to(dir_path))
                    if prefix:
                        arcname = f"{prefix}/{arcname}"
                    zf.write(file, arcname)
        buf.seek(0)

        zip_name = f"{dir_path.name}.zip"
        logger.debug("KB upload_dir: project=%s dir=%s (%d bytes)", self._slug, dir_path, buf.getbuffer().nbytes)
        data = KnowbaseProjectsDocumentsBulkCreateRequest(archive=(zip_name, buf, "application/zip"))
        result = self._api.projects_documents_bulk_create(project_slug=self._slug, data=data)
        logger.info("KB upload_dir queued: project=%s job_id=%s", self._slug, result.job_id)
        return result

    def job_status(self, job_id: str) -> ArchiveJobStatus:
        """Get archive ingestion job status.

        Args:
            job_id: Job ID returned by :meth:`upload_archive` or :meth:`upload_dir`.

        Returns:
            :class:`ArchiveJobStatus` with ``status`` ∈ {queued, processing, success, error}.
        """
        return self._api.projects_documents_bulk_status_retrieve(
            project_slug=self._slug, job_id=job_id
        )

    def wait_for_archive_done(
        self,
        job_id: str,
        *,
        timeout: float = 600.0,
        poll_interval: float = 3.0,
    ) -> ArchiveJobStatus:
        """Block until archive ingestion job completes (success or error).

        Args:
            job_id: Job ID returned by :meth:`upload_archive` or :meth:`upload_dir`.
            timeout: Max seconds to wait (default: 600).
            poll_interval: Seconds between status checks (default: 3).

        Returns:
            Final :class:`ArchiveJobStatus` with ``status='success'``.

        Raises:
            TimeoutError: If job doesn't complete within ``timeout`` seconds.
            RuntimeError: If job fails.
        """
        return self._poll_until_complete(
            job_uuid=job_id,
            get_status=lambda: self.job_status(job_id),
            get_result=lambda: self.job_status(job_id),
            status_attr="status",
            error_attr="error",
            success_status="success",
            error_status="error",
            poll_interval=poll_interval,
            timeout=timeout,
            error_message="Archive ingestion job failed",
        )

    def wait_for_processing(
        self,
        document_id: str,
        *,
        timeout: float = 120.0,
        poll_interval: float = 2.0,
    ) -> Document:
        """Block until a document's processing_status reaches 'completed' or 'failed'.

        Args:
            document_id: Document UUID.
            timeout: Max seconds to wait (default: 120).
            poll_interval: Seconds between polls (default: 2).

        Returns:
            Final :class:`Document` object with ``processing_status='completed'``.

        Raises:
            TimeoutError: If document doesn't complete within ``timeout`` seconds.
            RuntimeError: If processing fails.
        """
        start = time.time()
        while True:
            doc = self.get(document_id)
            status = str(doc.processing_status.value if hasattr(doc.processing_status, "value") else doc.processing_status)
            if status == "completed":
                return doc
            if status == "failed":
                error = getattr(doc, "processing_error", "") or "unknown error"
                raise RuntimeError(f"Document {document_id} processing failed: {error}")
            if time.time() - start >= timeout:
                raise TimeoutError(
                    f"Document {document_id} did not complete within {timeout}s (status: {status})"
                )
            time.sleep(poll_interval)


class _AsyncDocumentsResource(AsyncPollingMixin, AsyncBaseResource):
    """Documents within a KB project (async)."""

    def __init__(self, config: SDKConfig, project_slug: str) -> None:
        super().__init__(config, use_upload_timeout=True)
        self._api = _AsyncAPI(self._http_client)
        self._slug = project_slug

    @async_api_error_handler
    async def list(self, *, page: int = 1, page_size: int = 20) -> PaginatedDocumentList:
        """List documents in this project."""
        return await self._api.projects_documents_list(  # type: ignore[return-value]
            project_slug=self._slug, page=page, page_size=page_size,
        )

    @async_api_error_handler
    async def upload(
        self,
        *,
        title: str,
        content: str,
        source_url: str | None = None,
    ) -> Document:
        """Upload a document (Markdown/plain text)."""
        data = DocumentUploadRequest(title=title, content=content, source_url=source_url)
        return await self._api.projects_documents_create(project_slug=self._slug, data=data)

    @async_api_error_handler
    async def upload_file(self, path: Path, *, source_url: str | None = None) -> Document:
        """Upload a local Markdown or text file."""
        path = Path(path)
        content = path.read_text(encoding="utf-8")
        title = path.stem.replace("-", " ").replace("_", " ").title()
        return await self.upload(title=title, content=content, source_url=source_url)

    @async_api_error_handler
    async def get(self, document_id: str) -> Document:
        """Get a document by ID."""
        return await self._api.projects_documents_retrieve(id=document_id, project_slug=self._slug)

    @async_api_error_handler
    async def delete(self, document_id: str) -> None:
        """Delete a document and all its chunks."""
        await self._api.projects_documents_destroy(id=document_id, project_slug=self._slug)

    @async_api_error_handler
    async def upload_archive(self, zip_path: str | Path) -> ArchiveBulkCreateResponse:
        """Upload a ZIP archive for bulk document ingestion (async).

        Args:
            zip_path: Path to a ``.zip`` file on disk.

        Returns:
            :class:`ArchiveBulkCreateResponse` with ``job_id`` for polling.
        """
        zip_path = Path(zip_path)
        logger.debug("KB async upload_archive: project=%s path=%s", self._slug, zip_path)
        with zip_path.open("rb") as fh:
            data = KnowbaseProjectsDocumentsBulkCreateRequest(archive=(zip_path.name, fh, "application/zip"))
            result = await self._api.projects_documents_bulk_create(project_slug=self._slug, data=data)
        logger.info("KB async upload_archive queued: project=%s job_id=%s", self._slug, result.job_id)
        return result

    @async_api_error_handler
    async def upload_dir(self, dir_path: str | Path, *, prefix: str = "") -> ArchiveBulkCreateResponse:
        """Zip a local directory in-memory and upload it for bulk ingestion (async).

        Args:
            dir_path: Path to a local directory to zip and upload.
            prefix: Optional path prefix inside the ZIP (default: none).

        Returns:
            :class:`ArchiveBulkCreateResponse` with ``job_id`` for polling.
        """
        dir_path = Path(dir_path).resolve()
        if not dir_path.is_dir():
            raise ValueError(f"Not a directory: {dir_path}")

        buf = io.BytesIO()
        with zipfile.ZipFile(buf, "w", compression=zipfile.ZIP_DEFLATED) as zf:
            for file in sorted(dir_path.rglob("*")):
                if file.is_file():
                    arcname = str(file.relative_to(dir_path))
                    if prefix:
                        arcname = f"{prefix}/{arcname}"
                    zf.write(file, arcname)
        buf.seek(0)

        zip_name = f"{dir_path.name}.zip"
        logger.debug("KB async upload_dir: project=%s dir=%s (%d bytes)", self._slug, dir_path, buf.getbuffer().nbytes)
        data = KnowbaseProjectsDocumentsBulkCreateRequest(archive=(zip_name, buf, "application/zip"))
        result = await self._api.projects_documents_bulk_create(project_slug=self._slug, data=data)
        logger.info("KB async upload_dir queued: project=%s job_id=%s", self._slug, result.job_id)
        return result

    async def job_status(self, job_id: str) -> ArchiveJobStatus:
        """Get archive ingestion job status (async)."""
        return await self._api.projects_documents_bulk_status_retrieve(
            project_slug=self._slug, job_id=job_id
        )

    async def wait_for_archive_done(
        self,
        job_id: str,
        *,
        timeout: float = 600.0,
        poll_interval: float = 3.0,
    ) -> ArchiveJobStatus:
        """Block until archive ingestion job completes (async).

        Args:
            job_id: Job ID returned by :meth:`upload_archive` or :meth:`upload_dir`.
            timeout: Max seconds to wait (default: 600).
            poll_interval: Seconds between status checks (default: 3).

        Returns:
            Final :class:`ArchiveJobStatus` with ``status='success'``.

        Raises:
            TimeoutError: If job doesn't complete within ``timeout`` seconds.
            RuntimeError: If job fails.
        """
        return await self._poll_until_complete(
            job_uuid=job_id,
            get_status=lambda: self.job_status(job_id),
            get_result=lambda: self.job_status(job_id),
            status_attr="status",
            error_attr="error",
            success_status="success",
            error_status="error",
            poll_interval=poll_interval,
            timeout=timeout,
            error_message="Archive ingestion job failed",
        )

    async def wait_for_processing(
        self,
        document_id: str,
        *,
        timeout: float = 120.0,
        poll_interval: float = 2.0,
    ) -> Document:
        """Block until a document's processing_status reaches 'completed' or 'failed' (async).

        Args:
            document_id: Document UUID.
            timeout: Max seconds to wait (default: 120).
            poll_interval: Seconds between polls (default: 2).

        Returns:
            Final :class:`Document` with ``processing_status='completed'``.

        Raises:
            TimeoutError: If document doesn't complete within ``timeout`` seconds.
            RuntimeError: If processing fails.
        """
        start = time.time()
        while True:
            doc = await self.get(document_id)
            status = str(doc.processing_status.value if hasattr(doc.processing_status, "value") else doc.processing_status)
            if status == "completed":
                return doc
            if status == "failed":
                error = getattr(doc, "processing_error", "") or "unknown error"
                raise RuntimeError(f"Document {document_id} processing failed: {error}")
            if time.time() - start >= timeout:
                raise TimeoutError(
                    f"Document {document_id} did not complete within {timeout}s (status: {status})"
                )
            await asyncio.sleep(poll_interval)


__all__ = [
    "_DocumentsResource",
    "_AsyncDocumentsResource",
]
