Artifex/control_plane/authoring/sources.py

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass
from fnmatch import fnmatch
from pathlib import Path

from django.db import transaction
from django.db.models import Max

from control_plane.authoring.models import (
    DocumentAuthority,
    DocumentType,
    SourceDocument,
    SourceDocumentVersion,
    SourcePassage,
    Work,
)

SUPPORTED_SOURCE_SUFFIXES = {".json", ".log", ".md", ".txt"}


@dataclass(frozen=True)
class SourceRegistrationResult:
    path: Path
    logical_key: str
    status: str
    source_sha256: str
    version: int | None = None
    passage_count: int = 0


def discover_source_paths(root: Path, include_globs: list[str] | None = None) -> list[Path]:
    root = root.resolve()
    if root.is_file():
        supported = root.suffix.lower() in SUPPORTED_SOURCE_SUFFIXES
        included = not include_globs or any(
            fnmatch(root.name, pattern) for pattern in include_globs
        )
        return [root] if supported and included else []
    return sorted(
        path.resolve()
        for path in root.rglob("*")
        if path.is_file()
        and path.suffix.lower() in SUPPORTED_SOURCE_SUFFIXES
        and (
            not include_globs
            or any(fnmatch(path.relative_to(root).as_posix(), pattern) for pattern in include_globs)
        )
    )


def source_logical_key(path: Path, root: Path) -> str:
    path = path.resolve()
    root = root.resolve()
    if root.is_file():
        return path.name
    return path.relative_to(root).as_posix()


def source_title(path: Path, content: str) -> str:
    if path.suffix.lower() == ".md":
        match = re.search(r"^#{1,6}\s+(.+?)\s*$", content, flags=re.MULTILINE)
        if match:
            return match.group(1).strip()
    return path.stem.replace("-", " ").replace("_", " ").strip().title()


def passage_spans(content: str) -> list[dict[str, int | str]]:
    lines = content.splitlines(keepends=True)
    if not lines and content:
        lines = [content]
    passages: list[dict[str, int | str]] = []
    block_start_line: int | None = None
    block_start_char: int | None = None
    block_end_line = 0
    block_end_char = 0
    cursor = 0

    def finish_block() -> None:
        nonlocal block_start_line, block_start_char
        if block_start_line is None or block_start_char is None:
            return
        passage_content = content[block_start_char:block_end_char]
        passages.append(
            {
                "ordinal": len(passages) + 1,
                "start_line": block_start_line,
                "end_line": block_end_line,
                "start_char": block_start_char,
                "end_char": block_end_char,
                "content": passage_content,
                "sha256": hashlib.sha256(passage_content.encode("utf-8")).hexdigest(),
            }
        )
        block_start_line = None
        block_start_char = None

    for line_number, line in enumerate(lines, start=1):
        content_end = cursor + len(line.rstrip("\r\n"))
        if line.strip():
            if block_start_line is None:
                block_start_line = line_number
                block_start_char = cursor
            block_end_line = line_number
            block_end_char = content_end
        else:
            finish_block()
        cursor += len(line)
    finish_block()
    return passages


def inspect_source(path: Path, root: Path) -> SourceRegistrationResult:
    raw = path.read_bytes()
    content = raw.decode("utf-8")
    return SourceRegistrationResult(
        path=path.resolve(),
        logical_key=source_logical_key(path, root),
        status="discovered",
        source_sha256=hashlib.sha256(raw).hexdigest(),
        passage_count=len(passage_spans(content)),
    )


@transaction.atomic
def register_source(
    *,
    work: Work,
    path: Path,
    root: Path,
    authority: str,
    document_type: str = DocumentType.OTHER,
) -> SourceRegistrationResult:
    if authority not in DocumentAuthority.values:
        raise ValueError(f"unsupported document authority: {authority}")
    if document_type not in DocumentType.values:
        raise ValueError(f"unsupported document type: {document_type}")

    path = path.resolve()
    raw = path.read_bytes()
    content = raw.decode("utf-8")
    digest = hashlib.sha256(raw).hexdigest()
    logical_key = source_logical_key(path, root)
    document, _ = SourceDocument.objects.get_or_create(
        work=work,
        logical_key=logical_key,
        defaults={
            "title": source_title(path, content),
            "document_type": document_type,
        },
    )
    if document.document_type != document_type:
        raise ValueError(
            f"source {logical_key} is already registered as {document.document_type}, "
            f"not {document_type}"
        )
    latest = document.versions.order_by("-version").first()
    if (
        latest
        and latest.source_sha256 == digest
        and latest.authority == authority
        and latest.source_path == str(path)
    ):
        return SourceRegistrationResult(
            path=path,
            logical_key=logical_key,
            status="unchanged",
            source_sha256=digest,
            version=latest.version,
            passage_count=latest.passages.count(),
        )

    version_number = (document.versions.aggregate(value=Max("version"))["value"] or 0) + 1
    version = SourceDocumentVersion.objects.create(
        document=document,
        version=version_number,
        authority=authority,
        source_path=str(path),
        content=content,
        source_sha256=digest,
        byte_size=len(raw),
        supersedes=latest,
    )
    spans = passage_spans(content)
    SourcePassage.objects.bulk_create(
        [SourcePassage(document_version=version, **span) for span in spans]
    )
    return SourceRegistrationResult(
        path=path,
        logical_key=logical_key,
        status="created" if latest is None else "versioned",
        source_sha256=digest,
        version=version_number,
        passage_count=len(spans),
    )