"""Turning an uploaded object key into something a model can read.

An employee attaches a file by uploading it through the platform's
browser-direct presign flow and handing this domain the resulting S3 object
key. Three things then have to happen, in this order, and the order is the
security property:

1. **Verify the key.** ``parse_object_key`` recovers the purpose and the
   uploader id from the key's own shape. A key whose purpose is not an AI
   Labs purpose, or whose uploader is not the caller, is refused before any
   byte is read. Without this a crafted key would let an employee attach —
   and thereby read back through the model — someone else's evidence photo.
   This is the same gate the feed learned to apply after #2124.
2. **Read the bytes** from Helix's own bucket.
3. **Convert**, per kind:

   * images and audio become ``data:<type>;base64,...`` URIs;
   * documents are parsed **server-side** and only their extracted text
     travels to the vendor.

Why base64 rather than handing the vendor a presigned URL: verified on
2026-09-03 that DashScope accepts data URIs on all three surfaces that need
one (chat vision, image editing, video reference image). A presigned URL
would have worked too, but it means giving a third party a live link into
Helix's storage, and it makes every attachment depend on our bucket being
reachable from the vendor's network. Inlining costs ~1.33x in request size
and removes both.

The document text is **not** wrapped in instructions. It arrives as a text
part introduced by nothing more than the filename, because this product
promises not to author anything on the employee's behalf.
"""
from __future__ import annotations

import base64
import csv
import io
import logging
import uuid
from dataclasses import dataclass
from typing import Any

from sqlalchemy.ext.asyncio import AsyncSession

from app.core.exceptions import AppException
from app.core.storage import StorageNotConfiguredError, get_object_bytes
from app.core.storage_uploads import safe_display_filename
from app.domains.ai_labs import catalog
from app.storage_router import AI_LABS_PURPOSES, parse_object_key

logger = logging.getLogger(__name__)

KIND_IMAGE = "image"
KIND_DOCUMENT = "document"
KIND_AUDIO = "audio"

#: Purpose → kind. The purpose is baked into the object key at upload time,
#: so this is the authoritative classification: a client claiming a
#: different kind cannot change what the file actually is.
_PURPOSE_KIND = {
    "ai_labs_image": KIND_IMAGE,
    "ai_labs_document": KIND_DOCUMENT,
    "ai_labs_audio": KIND_AUDIO,
}

#: Rows read out of a spreadsheet. A workbook is a plausible thing to
#: attach and an implausible thing to send in full.
_XLSX_MAX_ROWS = 2000


@dataclass(frozen=True)
class Attachment:
    """One verified, fetched attachment, ready to be sent."""

    object_key: str
    kind: str
    content_type: str
    filename: str
    size_bytes: int
    #: Images / audio: the data URI. Documents: ``None``.
    data_uri: str | None = None
    #: Documents: extracted text. Others: ``None``.
    text: str | None = None
    #: True when the document was longer than the extraction ceiling. The
    #: UI states this; it is never a silent truncation.
    truncated: bool = False

    def as_stored(self) -> dict[str, Any]:
        """The persisted shape — a reference and metadata, never content."""
        return {
            "object_key": self.object_key,
            "kind": self.kind,
            "content_type": self.content_type,
            "filename": self.filename,
            "size_bytes": self.size_bytes,
            "truncated": self.truncated,
        }


def _refuse(message: str, *, code: str = "AI_LABS_ATTACHMENT_INVALID") -> AppException:
    return AppException(code=code, message=message, status_code=422)


def classify(object_key: str, *, owner_id: uuid.UUID) -> str:
    """Verify one key belongs to ``owner_id`` and return its kind.

    Raises 422 for a key this domain does not accept and 403 for someone
    else's key — distinguishable on purpose, because "wrong sort of file"
    and "not yours" call for different messages in the UI.
    """
    parsed = parse_object_key(object_key)
    if parsed is None:
        raise _refuse("Attachment key shape is not recognised.")
    purpose, uploader_id = parsed
    if purpose not in AI_LABS_PURPOSES:
        raise _refuse(
            "That file was uploaded for a different purpose and cannot be "
            "attached here.",
        )
    if str(uploader_id) != str(owner_id):
        raise AppException(
            code="AI_LABS_ATTACHMENT_NOT_OWNER",
            message="That file was uploaded by someone else.",
            status_code=403,
        )
    return _PURPOSE_KIND[purpose]


async def claim(
    db: AsyncSession,
    *,
    object_keys: list[str],
    owner_id: uuid.UUID,
    resource_type: str,
    resource_id: str,
) -> None:
    """Claim the upload grants for keys this domain is about to persist.

    Without this the grants stay ``status="issued"`` and the platform's
    hourly orphan sweep deletes the objects 24 hours later, leaving the
    employee's own row pointing at nothing. Every other domain that
    consumes ``/uploads/presign-put`` does the same thing.

    Wrapped rather than called directly because
    :func:`mark_uploads_referenced` verifies the landed object with a HEAD,
    which raises :class:`StorageNotConfiguredError` — **not** an
    ``AppException``, so it would surface as a 500 instead of the 503 this
    domain returns everywhere else storage is missing.

    Runs inside the caller's transaction on purpose: the claim and the row
    that references it have to commit together, or a crash between them
    recreates the orphan this exists to prevent.
    """
    if not object_keys:
        return
    from app.core.storage_uploads import mark_uploads_referenced

    try:
        await mark_uploads_referenced(
            db,
            object_keys=object_keys,
            person_id=owner_id,
            allowed_purposes=AI_LABS_PURPOSES,
            resource_type=resource_type,
            resource_id=resource_id,
        )
    except StorageNotConfiguredError as exc:
        raise AppException(
            code="STORAGE_NOT_CONFIGURED",
            message="File storage is not configured.",
            status_code=503,
        ) from exc


# ── document text extraction ───────────────────────────────────


def _extract_pdf(raw: bytes) -> str:
    from pypdf import PdfReader

    reader = PdfReader(io.BytesIO(raw))
    pieces: list[str] = []
    for page in reader.pages:
        try:
            pieces.append(page.extract_text() or "")
        except Exception as exc:  # noqa: BLE001 — one bad page, not a bad file
            # A single malformed page is common in real PDFs. Losing it is
            # much better than losing the document.
            logger.debug("ai_labs: skipping unreadable PDF page: %s", exc)
    return "\n\n".join(piece for piece in pieces if piece.strip())


def _extract_docx(raw: bytes) -> str:
    import docx

    document = docx.Document(io.BytesIO(raw))
    pieces = [para.text for para in document.paragraphs if para.text.strip()]
    for table in document.tables:
        for row in table.rows:
            cells = [cell.text.strip() for cell in row.cells]
            if any(cells):
                pieces.append("\t".join(cells))
    return "\n".join(pieces)


def _extract_xlsx(raw: bytes) -> str:
    from openpyxl import load_workbook

    # ``read_only`` streams rows instead of building the whole object graph,
    # and ``data_only`` reads cached formula results rather than "=A1*B1".
    workbook = load_workbook(
        io.BytesIO(raw), read_only=True, data_only=True,
    )
    pieces: list[str] = []
    rows_seen = 0
    try:
        for sheet in workbook.worksheets:
            pieces.append(f"# {sheet.title}")
            for row in sheet.iter_rows(values_only=True):
                if rows_seen >= _XLSX_MAX_ROWS:
                    pieces.append("…")
                    return "\n".join(pieces)
                rows_seen += 1
                cells = [
                    "" if value is None else str(value) for value in row
                ]
                if any(cell.strip() for cell in cells):
                    pieces.append("\t".join(cells))
    finally:
        workbook.close()
    return "\n".join(pieces)


def _extract_csv(raw: bytes) -> str:
    text = raw.decode("utf-8", errors="replace")
    reader = csv.reader(io.StringIO(text))
    rows = []
    for index, row in enumerate(reader):
        if index >= _XLSX_MAX_ROWS:
            rows.append("…")
            break
        rows.append("\t".join(row))
    return "\n".join(rows)


def extract_document_text(raw: bytes, content_type: str) -> tuple[str, bool]:
    """``(text, truncated)`` for a supported document type.

    Extraction failures raise 422 with the filename-free reason, because
    "we could not read this file" is actionable for the employee in a way
    that a 500 is not.
    """
    try:
        if content_type == "application/pdf":
            text = _extract_pdf(raw)
        elif content_type.endswith(".wordprocessingml.document"):
            text = _extract_docx(raw)
        elif content_type.endswith(".spreadsheetml.sheet"):
            text = _extract_xlsx(raw)
        elif content_type == "text/csv":
            text = _extract_csv(raw)
        else:
            text = raw.decode("utf-8", errors="replace")
    except AppException:
        raise
    except Exception as exc:  # noqa: BLE001 — vendor-quality files vary wildly
        logger.info(
            "ai_labs: document extraction failed for %s: %s", content_type, exc,
        )
        raise _refuse(
            "This document could not be read. It may be encrypted, "
            "corrupt, or an image-only scan with no text layer.",
            code="AI_LABS_DOCUMENT_UNREADABLE",
        ) from exc

    text = text.strip()
    if not text:
        raise _refuse(
            "No text could be extracted from this document — an image-only "
            "scan has nothing for the model to read.",
            code="AI_LABS_DOCUMENT_EMPTY",
        )
    if len(text) > catalog.DOCUMENT_TEXT_MAX_CHARS:
        return text[: catalog.DOCUMENT_TEXT_MAX_CHARS], True
    return text, False


# ── loading ────────────────────────────────────────────────────


async def load(
    db: AsyncSession,
    *,
    object_keys: list[str],
    owner_id: uuid.UUID,
    allowed_kinds: frozenset[str] | set[str],
    names: dict[str, str] | None = None,
) -> list[Attachment]:
    """Verify, fetch and convert every key. Order is preserved.

    ``allowed_kinds`` lets each surface state what it can actually use — a
    video reference image accepts an image and nothing else, so attaching a
    PDF there fails with a clear message rather than being ignored.

    ``names`` maps object key to the display filename the caller wants
    shown. Untrusted input, so it is sanitised and used **only** as a
    label: every content decision (kind, content type) still comes from the
    stored key.
    """
    if not object_keys:
        return []
    if len(object_keys) > catalog.MAX_ATTACHMENTS_PER_TURN:
        raise _refuse(
            f"At most {catalog.MAX_ATTACHMENTS_PER_TURN} attachments per "
            "message.",
        )

    loaded: list[Attachment] = []
    for object_key in object_keys:
        kind = classify(object_key, owner_id=owner_id)
        if kind not in allowed_kinds:
            raise _refuse(
                f"A {kind} attachment cannot be used here.",
            )
        try:
            raw = await get_object_bytes(db, object_key=object_key)
        except StorageNotConfiguredError as exc:
            raise AppException(
                code="STORAGE_NOT_CONFIGURED",
                message="File storage is not configured.",
                status_code=503,
            ) from exc
        if raw is None:
            raise _refuse(
                "That attachment is no longer in storage. Please upload it "
                "again.",
                code="AI_LABS_ATTACHMENT_MISSING",
            )
        if len(raw) > catalog.MAX_ATTACHMENT_BYTES:
            # The presign flow already caps size per purpose; this is the
            # server-side backstop for an object that grew or was written
            # by some other path.
            raise _refuse(
                "That attachment is larger than this surface accepts.",
                code="AI_LABS_ATTACHMENT_TOO_LARGE",
            )

        # The display name comes from the caller when it has one, because
        # AI Labs keys deliberately do NOT carry it (see
        # ``storage_uploads.OPAQUE_NAME_PURPOSES``: the key is copied into
        # the grant ledger and the company-readable audit row, and the
        # employee's filename is their content). The key-derived fallback
        # keeps this function usable where no name was sent — generation
        # inputs, whose filenames are never displayed.
        filename = (names or {}).get(object_key) or ""
        if not filename:
            filename = object_key.rsplit("/", 1)[-1]
            # Legacy layout is "<random12>-<safe filename>"; the random
            # prefix is plumbing and showing it would be noise.
            if "-" in filename:
                filename = filename.split("-", 1)[1] or filename
        filename = safe_display_filename(filename)
        extension_source = object_key.rsplit("/", 1)[-1]

        if kind == KIND_DOCUMENT:
            content_type = _document_type_from_name(extension_source)
            text, truncated = extract_document_text(raw, content_type)
            loaded.append(Attachment(
                object_key=object_key, kind=kind, content_type=content_type,
                filename=filename, size_bytes=len(raw),
                text=text, truncated=truncated,
            ))
            continue

        content_type = (
            _image_type_from_name(extension_source)
            if kind == KIND_IMAGE
            else _audio_type_from_name(extension_source)
        )
        encoded = base64.b64encode(raw).decode("ascii")
        loaded.append(Attachment(
            object_key=object_key, kind=kind, content_type=content_type,
            filename=filename, size_bytes=len(raw),
            data_uri=f"data:{content_type};base64,{encoded}",
        ))
    return loaded


#: Extension → content type. The upload grant recorded the declared type,
#: but re-reading that row per attachment is a query per file for a value
#: the key's own suffix already carries; the presign endpoint validated the
#: type against the purpose's allow-list at upload time, so the extension
#: and the declared type agree by construction.
_IMAGE_TYPES = {
    "jpg": "image/jpeg", "jpeg": "image/jpeg",
    "png": "image/png", "webp": "image/webp",
}
_DOCUMENT_TYPES = {
    "pdf": "application/pdf",
    "docx": (
        "application/vnd.openxmlformats-officedocument"
        ".wordprocessingml.document"
    ),
    "xlsx": "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
    "txt": "text/plain",
    "md": "text/markdown",
    "csv": "text/csv",
}
_AUDIO_TYPES = {
    "mp3": "audio/mpeg", "m4a": "audio/mp4", "mp4": "audio/mp4",
    "wav": "audio/wav", "webm": "audio/webm",
    "aac": "audio/aac", "ogg": "audio/ogg",
}


def _suffix(filename: str) -> str:
    return filename.rsplit(".", 1)[-1].lower() if "." in filename else ""


def _image_type_from_name(filename: str) -> str:
    kind = _IMAGE_TYPES.get(_suffix(filename))
    if kind is None:
        raise _refuse("Unsupported image type.")
    return kind


def _document_type_from_name(filename: str) -> str:
    kind = _DOCUMENT_TYPES.get(_suffix(filename))
    if kind is None:
        raise _refuse("Unsupported document type.")
    return kind


def _audio_type_from_name(filename: str) -> str:
    kind = _AUDIO_TYPES.get(_suffix(filename))
    if kind is None:
        raise _refuse("Unsupported audio type.")
    return kind


__all__ = [
    "KIND_AUDIO",
    "KIND_DOCUMENT",
    "KIND_IMAGE",
    "Attachment",
    "classify",
    "extract_document_text",
    "load",
]