import hashlib
import mimetypes
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING
from urllib.parse import urlparse

import httpx

from haiku.rag.client.processing import ensure_chunks_embedded
from haiku.rag.client.titles import resolve_title
from haiku.rag.converters import get_converter
from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document
from haiku.rag.store.models.document_item import extract_items

if TYPE_CHECKING:
    from docling_core.types.doc.document import DoclingDocument

    from haiku.rag.client import HaikuRAG


async def _store_document_with_chunks(
    client: "HaikuRAG",
    document: Document,
    chunks: list[Chunk],
    docling_document: "DoclingDocument",
) -> Document:
    """Store a document with chunks, embedding any that lack embeddings.

    Handles versioning/rollback on failure.
    """
    # Ensure all chunks have embeddings before storing
    chunks = await ensure_chunks_embedded(client._config, chunks)

    # Snapshot table versions for versioned rollback (if supported)
    versions = await client.store.current_table_versions()

    # Create the document
    created_doc = await client.document_repository.create(document)

    try:
        assert created_doc.id is not None, (
            "Document ID should not be None after creation"
        )
        # Set document_id and order for all chunks
        for order, chunk in enumerate(chunks):
            chunk.document_id = created_doc.id
            chunk.order = order

        # Batch create all chunks in a single operation
        await client.chunk_repository.create(chunks)

        # Extract and store document items for context expansion
        items = extract_items(created_doc.id, docling_document)
        await client.document_item_repository.create_items(created_doc.id, items)

        # Vacuum old versions in background (non-blocking) if auto_vacuum enabled
        if client._config.storage.auto_vacuum:
            client._schedule_vacuum()

        return created_doc
    except Exception:
        # Roll back to the captured versions and re-raise
        await client.store.restore_table_versions(versions)
        raise


async def _update_document_with_chunks(
    client: "HaikuRAG",
    document: Document,
    chunks: list[Chunk],
    docling_document: "DoclingDocument | None" = None,
) -> Document:
    """Update a document and replace its chunks, embedding any that lack embeddings.

    Handles versioning/rollback on failure. When `docling_document` is None,
    existing items are preserved.
    """
    assert document.id is not None, "Document ID is required for update"

    # Snapshot existing picture bytes before deleting items so the post-delete
    # extract_items can merge them back when the live docling has had its
    # picture URIs stripped (rebuild / re-extract via the stored blob).
    existing_picture_data: dict[str, bytes] | None = None
    if docling_document is not None:
        existing_picture_data = (
            await client.document_item_repository.get_all_picture_data(document.id)
        )

    chunks = await ensure_chunks_embedded(client._config, chunks)

    versions = await client.store.current_table_versions()

    # Delete existing chunks before writing new ones
    await client.chunk_repository.delete_by_document_id(document.id)

    try:
        updated_doc = await client.document_repository.update(document)

        assert updated_doc.id is not None
        for order, chunk in enumerate(chunks):
            chunk.document_id = updated_doc.id
            chunk.order = order

        await client.chunk_repository.create(chunks)

        if docling_document is not None:
            await client.document_item_repository.delete_by_document_id(updated_doc.id)
            items = extract_items(
                updated_doc.id,
                docling_document,
                existing_picture_data=existing_picture_data,
            )
            await client.document_item_repository.create_items(updated_doc.id, items)

        if client._config.storage.auto_vacuum:
            client._schedule_vacuum()

        return updated_doc
    except Exception:
        await client.store.restore_table_versions(versions)
        raise


async def create_document(
    client: "HaikuRAG",
    content: str,
    uri: str | None = None,
    title: str | None = None,
    metadata: dict | None = None,
    format: str = "md",
) -> Document:
    """Create a new document from text content.

    Converts the content, chunks it, and generates embeddings.
    """
    from haiku.rag.embeddings import embed_chunks

    # Convert → Chunk → Embed using primitives
    converter = get_converter(client._config)
    docling_document = await converter.convert_text(content, format=format)
    chunks = await client.chunk(docling_document)
    embedded_chunks = await embed_chunks(chunks, client._config)

    # Store markdown export as content for better display/readability.
    # The original is preserved in docling_document.
    stored_content = docling_document.export_to_markdown()

    if title is None:
        title = await resolve_title(client._config, docling_document, stored_content)

    document = Document(
        content=stored_content,
        uri=uri,
        title=title,
        metadata=metadata or {},
    )
    document.set_docling(docling_document)

    return await _store_document_with_chunks(
        client, document, embedded_chunks, docling_document
    )


async def import_document(
    client: "HaikuRAG",
    docling_document: "DoclingDocument",
    chunks: list[Chunk],
    uri: str | None = None,
    title: str | None = None,
    metadata: dict | None = None,
) -> Document:
    """Import a pre-processed document with chunks.

    Use this when conversion, chunking, and embedding were done externally.
    Chunks without embeddings will be automatically embedded.
    """
    content = docling_document.export_to_markdown()
    if title is None:
        title = await resolve_title(client._config, docling_document, content)

    document = Document(
        content=content,
        uri=uri,
        title=title,
        metadata=metadata or {},
    )
    document.set_docling(docling_document)

    return await _store_document_with_chunks(client, document, chunks, docling_document)


async def create_document_from_source(
    client: "HaikuRAG",
    source: str | Path,
    title: str | None = None,
    metadata: dict | None = None,
    uri: str | None = None,
    storage_options: dict[str, str] | None = None,
) -> Document | list[Document]:
    """Create or update document(s) from a file path, directory, or URL.

    Checks if a document with the same URI already exists:
    - If MD5 is unchanged, returns existing document
    - If MD5 changed, updates the document
    - If no document exists, creates a new one

    If ``uri`` is provided, it overrides the URI auto-derived from the source
    (which is normally ``file://`` for local files or the URL for remote
    sources). This is useful when callers want to persist documents under a
    logical identifier (e.g. an ArXiv ID) rather than the on-disk path. Not
    supported for directory sources, which produce one document per file.

    Returns a single Document for files/URLs, a list for directories.
    """
    metadata = metadata or {}

    source_str = str(source)
    parsed_url = urlparse(source_str)
    if parsed_url.scheme in ("http", "https"):
        return await _create_or_update_document_from_url(
            client, source_str, title=title, metadata=metadata, uri=uri
        )
    elif parsed_url.scheme == "s3":
        return await _create_or_update_document_from_s3(
            client,
            source_str,
            storage_options=storage_options,
            title=title,
            metadata=metadata,
            uri=uri,
        )
    elif parsed_url.scheme == "file":
        source_path = Path(parsed_url.path)
    else:
        source_path = Path(source) if isinstance(source, str) else source

    if source_path.is_dir():
        if uri is not None:
            raise ValueError(
                "uri override is not supported for directory sources; each file "
                "produces its own document with its own auto-derived URI."
            )
        from haiku.rag.monitor import FileFilter

        documents = []
        filter = FileFilter(
            ignore_patterns=client._config.monitor.ignore_patterns or None,
            include_patterns=client._config.monitor.include_patterns or None,
        )
        for path in source_path.rglob("*"):
            if path.is_file() and filter.include_file(str(path)):
                doc = await _create_document_from_file(
                    client, path, title=None, metadata=metadata
                )
                documents.append(doc)
        return documents

    return await _create_document_from_file(
        client, source_path, title=title, metadata=metadata, uri=uri
    )


async def _create_document_from_file(
    client: "HaikuRAG",
    source_path: Path,
    title: str | None = None,
    metadata: dict | None = None,
    uri: str | None = None,
) -> Document:
    """Create or update a document from a single file path.

    ``uri`` overrides the auto-derived ``file://`` URI; it's used as the
    canonical document identifier for lookup and storage.
    """
    from haiku.rag.embeddings import embed_chunks

    metadata = metadata or {}

    converter = get_converter(client._config)
    if source_path.suffix.lower() not in converter.supported_extensions:
        raise ValueError(f"Unsupported file extension: {source_path.suffix}")

    if not source_path.exists():
        raise ValueError(f"File does not exist: {source_path}")

    if uri is None:
        uri = source_path.absolute().as_uri()
    md5_hash = hashlib.md5(source_path.read_bytes(), usedforsecurity=False).hexdigest()

    content_type, _ = mimetypes.guess_type(str(source_path))
    if not content_type:
        content_type = "application/octet-stream"
    metadata.update({"contentType": content_type, "md5": md5_hash})

    # Check if document already exists
    existing_doc = await client.get_document_by_uri(uri)
    if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
        # MD5 unchanged; update title/metadata if provided
        updated = False
        if title is not None and title != existing_doc.title:
            existing_doc.title = title
            updated = True

        merged_metadata = {**(existing_doc.metadata or {}), **metadata}
        if merged_metadata != existing_doc.metadata:
            existing_doc.metadata = merged_metadata
            updated = True

        if updated:
            return await client.document_repository.update(existing_doc)
        return existing_doc

    # Convert → Chunk → Embed
    docling_document = await client.convert(source_path)
    chunks = await client.chunk(docling_document)
    embedded_chunks = await embed_chunks(chunks, client._config)

    stored_content = docling_document.export_to_markdown()

    if existing_doc:
        # Update existing document and rechunk
        existing_doc.content = stored_content
        existing_doc.metadata = metadata
        existing_doc.set_docling(docling_document)
        if title is not None:
            existing_doc.title = title
        elif existing_doc.title is None:
            existing_doc.title = await resolve_title(
                client._config, docling_document, stored_content
            )
        return await _update_document_with_chunks(
            client, existing_doc, embedded_chunks, docling_document
        )
    else:
        if title is None:
            title = await resolve_title(
                client._config, docling_document, stored_content
            )
        document = Document(
            content=stored_content,
            uri=uri,
            title=title,
            metadata=metadata,
        )
        document.set_docling(docling_document)
        return await _store_document_with_chunks(
            client, document, embedded_chunks, docling_document
        )


async def _create_or_update_document_from_url(
    client: "HaikuRAG",
    url: str,
    title: str | None = None,
    metadata: dict | None = None,
    uri: str | None = None,
) -> Document:
    """Create or update a document from a URL by downloading and parsing the content.

    ``uri`` overrides the URL as the stored document identifier.
    """
    from haiku.rag.client.processing import get_extension_from_content_type_or_url
    from haiku.rag.embeddings import embed_chunks

    metadata = metadata or {}
    stored_uri = uri if uri is not None else url

    converter = get_converter(client._config)
    supported_extensions = converter.supported_extensions

    async with httpx.AsyncClient() as http:
        response = await http.get(url)
        response.raise_for_status()

        md5_hash = hashlib.md5(response.content).hexdigest()

        content_type = response.headers.get("content-type", "").lower()

        # Check if document already exists
        existing_doc = await client.get_document_by_uri(stored_uri)
        if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
            updated = False
            if title is not None and title != existing_doc.title:
                existing_doc.title = title
                updated = True

            metadata.update({"contentType": content_type, "md5": md5_hash})
            merged_metadata = {**(existing_doc.metadata or {}), **metadata}
            if merged_metadata != existing_doc.metadata:
                existing_doc.metadata = merged_metadata
                updated = True

            if updated:
                return await client.document_repository.update(existing_doc)
            return existing_doc

        file_extension = get_extension_from_content_type_or_url(url, content_type)

        if file_extension not in supported_extensions:
            raise ValueError(
                f"Unsupported content type/extension: {content_type}/{file_extension}"
            )

        with tempfile.NamedTemporaryFile(
            mode="wb", suffix=file_extension, delete=False
        ) as temp_file:
            temp_file.write(response.content)
            temp_file.flush()
            temp_path = Path(temp_file.name)

        try:
            docling_document = await client.convert(temp_path, source_uri=url)
            chunks = await client.chunk(docling_document)
            embedded_chunks = await embed_chunks(chunks, client._config)
        finally:
            temp_path.unlink(missing_ok=True)

        metadata.update({"contentType": content_type, "md5": md5_hash})

        stored_content = docling_document.export_to_markdown()

        if existing_doc:
            existing_doc.content = stored_content
            existing_doc.metadata = metadata
            existing_doc.set_docling(docling_document)
            if title is not None:
                existing_doc.title = title
            elif existing_doc.title is None:
                existing_doc.title = await resolve_title(
                    client._config, docling_document, stored_content
                )
            return await _update_document_with_chunks(
                client, existing_doc, embedded_chunks, docling_document
            )
        else:
            if title is None:
                title = await resolve_title(
                    client._config, docling_document, stored_content
                )
            document = Document(
                content=stored_content,
                uri=stored_uri,
                title=title,
                metadata=metadata,
            )
            document.set_docling(docling_document)
            return await _store_document_with_chunks(
                client, document, embedded_chunks, docling_document
            )


async def _create_or_update_document_from_s3(
    client: "HaikuRAG",
    url: str,
    *,
    storage_options: dict[str, str] | None = None,
    title: str | None = None,
    metadata: dict | None = None,
    uri: str | None = None,
) -> Document:
    """Create or update a document from an s3:// URL.

    Two-stage change detection:
    - HEAD ETag matches stored metadata["etag"] → skip GET and re-chunk.
    - ETag differs but content MD5 matches → refresh etag only, no re-chunk.
    - Otherwise download, convert, chunk, embed.

    ``uri`` overrides the s3:// URL as the stored document identifier.
    """
    import obstore  # type: ignore[import-not-found]

    from haiku.rag.client.processing import get_extension_from_content_type_or_url
    from haiku.rag.embeddings import embed_chunks
    from haiku.rag.s3 import make_s3_store

    metadata = metadata or {}
    stored_uri = uri if uri is not None else url

    parsed = urlparse(url)
    bucket = parsed.netloc
    key = parsed.path.lstrip("/")
    if not bucket or not key:
        raise ValueError(f"Invalid S3 URI: {url}")

    converter = get_converter(client._config)
    supported_extensions = converter.supported_extensions

    store = make_s3_store(bucket, storage_options)

    head = await obstore.head_async(store, key)
    etag = (head.get("e_tag") or "").strip('"')

    existing_doc = await client.get_document_by_uri(stored_uri)
    if existing_doc and existing_doc.metadata.get("etag") == etag:
        updated = False
        if title is not None and title != existing_doc.title:
            existing_doc.title = title
            updated = True

        merged_metadata = {**(existing_doc.metadata or {}), **metadata}
        if merged_metadata != existing_doc.metadata:
            existing_doc.metadata = merged_metadata
            updated = True

        if updated:
            return await client.document_repository.update(existing_doc)
        return existing_doc

    content_type, _ = mimetypes.guess_type(key)
    if not content_type:
        content_type = "application/octet-stream"

    file_extension = get_extension_from_content_type_or_url(url, content_type)
    if file_extension not in supported_extensions:
        raise ValueError(
            f"Unsupported content type/extension: {content_type}/{file_extension}"
        )

    get_resp = await obstore.get_async(store, key)
    body = await get_resp.bytes_async()

    md5_hash = hashlib.md5(body, usedforsecurity=False).hexdigest()

    with tempfile.NamedTemporaryFile(
        mode="wb", suffix=file_extension, delete=False
    ) as temp_file:
        temp_file.write(body)
        temp_file.flush()
        temp_path = Path(temp_file.name)

    metadata.update({"contentType": content_type, "md5": md5_hash, "etag": etag})

    if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
        temp_path.unlink(missing_ok=True)
        merged_metadata = {**(existing_doc.metadata or {}), **metadata}
        updated = False
        if merged_metadata != existing_doc.metadata:
            existing_doc.metadata = merged_metadata
            updated = True
        if title is not None and title != existing_doc.title:
            existing_doc.title = title
            updated = True
        if updated:
            return await client.document_repository.update(existing_doc)
        return existing_doc

    try:
        docling_document = await client.convert(temp_path, source_uri=url)
        chunks = await client.chunk(docling_document)
        embedded_chunks = await embed_chunks(chunks, client._config)
    finally:
        temp_path.unlink(missing_ok=True)

    stored_content = docling_document.export_to_markdown()

    if existing_doc:
        existing_doc.content = stored_content
        existing_doc.metadata = metadata
        existing_doc.set_docling(docling_document)
        if title is not None:
            existing_doc.title = title
        elif existing_doc.title is None:
            existing_doc.title = await resolve_title(
                client._config, docling_document, stored_content
            )
        return await _update_document_with_chunks(
            client, existing_doc, embedded_chunks, docling_document
        )
    else:
        if title is None:
            title = await resolve_title(
                client._config, docling_document, stored_content
            )
        document = Document(
            content=stored_content,
            uri=stored_uri,
            title=title,
            metadata=metadata,
        )
        document.set_docling(docling_document)
        return await _store_document_with_chunks(
            client, document, embedded_chunks, docling_document
        )


async def update_document(
    client: "HaikuRAG",
    document_id: str,
    content: str | None = None,
    metadata: dict | None = None,
    chunks: list[Chunk] | None = None,
    title: str | None = None,
    docling_document: "DoclingDocument | None" = None,
) -> Document:
    """Update a document by ID.

    Updates specified fields. When content or docling_document is provided, the
    document is rechunked and re-embedded. Updates to only metadata or title
    skip rechunking for efficiency.

    Raises:
        ValueError: If document not found, or if both content and
            docling_document are provided.
    """
    from haiku.rag.embeddings import embed_chunks

    # Validate: content and docling_document are mutually exclusive
    if content is not None and docling_document is not None:
        raise ValueError(
            "content and docling_document are mutually exclusive. "
            "Provide one or the other, not both."
        )

    existing_doc = await client.get_document_by_id(document_id)
    if existing_doc is None:
        raise ValueError(f"Document with ID {document_id} not found")

    if title is not None:
        existing_doc.title = title
    if metadata is not None:
        existing_doc.metadata = metadata

    # Only metadata/title update - no rechunking needed
    if content is None and chunks is None and docling_document is None:
        return await client.document_repository.update(existing_doc)

    # Custom chunks provided - use them as-is
    if chunks is not None:
        if docling_document is not None:
            existing_doc.content = docling_document.export_to_markdown()
            existing_doc.set_docling(docling_document)
        elif content is not None:
            existing_doc.content = content

        return await _update_document_with_chunks(
            client, existing_doc, chunks, docling_document
        )

    # DoclingDocument provided without chunks - chunk and embed
    if docling_document is not None:
        existing_doc.content = docling_document.export_to_markdown()
        existing_doc.set_docling(docling_document)

        new_chunks = await client.chunk(docling_document)
        embedded_chunks = await embed_chunks(new_chunks, client._config)
        return await _update_document_with_chunks(
            client, existing_doc, embedded_chunks, docling_document
        )

    # Content provided without chunks - convert, chunk, and embed
    assert content is not None
    existing_doc.content = content
    converter = get_converter(client._config)
    converted_docling = await converter.convert_text(existing_doc.content, format="md")
    existing_doc.set_docling(converted_docling)

    new_chunks = await client.chunk(converted_docling)
    embedded_chunks = await embed_chunks(new_chunks, client._config)
    return await _update_document_with_chunks(
        client, existing_doc, embedded_chunks, converted_docling
    )


def check_source_accessible(uri: str) -> bool:
    """Check if a document's source URI is accessible."""
    parsed_url = urlparse(uri)
    try:
        if parsed_url.scheme == "file":
            return Path(parsed_url.path).exists()
        elif parsed_url.scheme in ("http", "https", "s3"):
            return True
        return False
    except Exception:
        return False