Agent-System-Backend/internal/store/document_delivery.go

package store

import (
	"crypto/sha256"
	"database/sql"
	"encoding/hex"
	"errors"
	"fmt"
	"regexp"
	"strings"
	"time"

	core "agent-system-core"
)

type DocumentDeliveryState string

const (
	DocumentDeliveryFrozen      DocumentDeliveryState = "frozen"
	DocumentDeliverySourceReady DocumentDeliveryState = "source-ready"
	DocumentDeliveryFailed      DocumentDeliveryState = "failed"
	DocumentDeliveryComplete    DocumentDeliveryState = "complete"
)

type DocumentDeliveryArtifactRole string

const (
	DocumentArtifactMarkdown DocumentDeliveryArtifactRole = "markdown"
	DocumentArtifactDOCX     DocumentDeliveryArtifactRole = "docx"
	DocumentArtifactManifest DocumentDeliveryArtifactRole = "manifest"
	DocumentArtifactArchive  DocumentDeliveryArtifactRole = "archive"
)

const (
	DocumentDeliveryFaultAfterMarkdownBlob    = "after-markdown-blob"
	DocumentDeliveryFaultBeforeMarkdownCommit = "before-markdown-commit"
	DocumentDeliveryFaultAfterDOCXBlob        = "after-docx-blob"
	DocumentDeliveryFaultAfterManifestBlob    = "after-manifest-blob"
	DocumentDeliveryFaultAfterArchiveBlob     = "after-archive-blob"
	DocumentDeliveryFaultBeforeDerivedCommit  = "before-derived-commit"
	DocumentDeliveryFaultAfterDerivedCommit   = "after-derived-commit"
)

var (
	ErrDocumentDeliveryConflict       = errors.New("document delivery operation conflicts with its frozen request")
	ErrDocumentDeliveryOutcomeUnknown = errors.New("document delivery commit outcome is unknown")
	deliveryOperationPattern          = regexp.MustCompile(`^[A-Za-z0-9][A-Za-z0-9._:-]{0,127}$`)
)

type DocumentDeliveryBeginInput struct {
	OperationID      string
	RequestSHA256    string
	ProjectID        string
	ConversationID   string
	SourceRun        string
	DocumentID       string
	Mode             string
	RendererIdentity string
}

type DocumentDeliveryArtifactInput struct {
	Role    DocumentDeliveryArtifactRole
	Name    string
	Kind    string
	MIME    string
	SHA256  string
	Content []byte
}

type DocumentDeliveryArtifact struct {
	Role     DocumentDeliveryArtifactRole `json:"role"`
	Artifact core.Artifact                `json:"artifact"`
	SHA256   string                       `json:"sha256"`
}

type DocumentDeliveryRecord struct {
	OperationID        string                     `json:"operationId"`
	Owner              string                     `json:"owner"`
	ProjectID          string                     `json:"projectId"`
	ConversationID     string                     `json:"conversationId"`
	SourceRun          string                     `json:"sourceRun,omitempty"`
	RequestSHA256      string                     `json:"requestSha256"`
	State              DocumentDeliveryState      `json:"state"`
	FrozenAt           time.Time                  `json:"frozenAt"`
	RendererIdentity   string                     `json:"rendererIdentity"`
	Mode               string                     `json:"mode"`
	DocumentID         string                     `json:"documentId"`
	CanonicalHash      string                     `json:"canonicalHash,omitempty"`
	ProfileHash        string                     `json:"profileHash,omitempty"`
	ReviewHash         string                     `json:"reviewHash,omitempty"`
	DeliveryKey        string                     `json:"deliveryKey,omitempty"`
	MarkdownArtifactID string                     `json:"markdownArtifactId,omitempty"`
	Error              string                     `json:"error,omitempty"`
	Artifacts          []DocumentDeliveryArtifact `json:"artifacts,omitempty"`
}

type storedDocumentDelivery struct {
	DocumentDeliveryRecord
	markdownSHA, docxID, docxSHA, manifestID, manifestSHA, archiveID, archiveSHA string
}

const documentDeliveryColumns = `operation_id,owner,project_id,conversation_id,source_run,request_sha256,state,frozen_at,renderer_identity,mode,document_id,canonical_hash,profile_hash,review_hash,delivery_key,markdown_artifact_id,markdown_sha256,docx_artifact_id,docx_sha256,manifest_artifact_id,manifest_sha256,archive_artifact_id,archive_sha256,error_text`

func (s *Store) BeginDocumentDeliveryAccountChecked(actor string, input DocumentDeliveryBeginInput, allowImplicitDefault bool) (DocumentDeliveryRecord, error) {
	if err := validateDocumentDeliveryBegin(actor, input); err != nil {
		return DocumentDeliveryRecord{}, err
	}
	var record DocumentDeliveryRecord
	err := s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		if err := s.cardResourcesForWrite(tx, actor, input.ProjectID, input.ConversationID); err != nil {
			return err
		}
		stored, found, err := s.documentDeliveryForWrite(tx, actor, input.OperationID)
		if err != nil {
			return err
		}
		if found {
			if !sameDocumentDeliveryInput(stored.DocumentDeliveryRecord, input) {
				return ErrDocumentDeliveryConflict
			}
			record = stored.DocumentDeliveryRecord
			return nil
		}
		now := time.Now().UTC()
		_, err = tx.Exec(s.q(`INSERT INTO document_render_deliveries(operation_id,owner,project_id,conversation_id,source_run,request_sha256,state,frozen_at,renderer_identity,mode,document_id,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?,?)`),
			input.OperationID, actor, input.ProjectID, input.ConversationID, input.SourceRun, input.RequestSHA256,
			DocumentDeliveryFrozen, now.Format(time.RFC3339Nano), input.RendererIdentity, input.Mode, input.DocumentID,
			now.Format(time.RFC3339Nano), now.Format(time.RFC3339Nano))
		if err == nil {
			record = DocumentDeliveryRecord{OperationID: input.OperationID, Owner: actor, ProjectID: input.ProjectID,
				ConversationID: input.ConversationID, SourceRun: input.SourceRun, RequestSHA256: input.RequestSHA256,
				State: DocumentDeliveryFrozen, FrozenAt: now, RendererIdentity: input.RendererIdentity, Mode: input.Mode, DocumentID: input.DocumentID}
		}
		return err
	})
	if err != nil {
		return record, err
	}
	return s.GetDocumentDelivery(actor, input.OperationID)
}

func (s *Store) PersistDocumentDeliveryMarkdownAccountChecked(actor, operationID string, artifact DocumentDeliveryArtifactInput, canonicalHash string, allowImplicitDefault bool) (DocumentDeliveryRecord, error) {
	if artifact.Role != DocumentArtifactMarkdown || canonicalHash == "" {
		return DocumentDeliveryRecord{}, errors.New("document delivery Markdown role and canonical hash are required")
	}
	if err := validateDeliveryArtifactInput(artifact); err != nil {
		return DocumentDeliveryRecord{}, err
	}
	current, err := s.GetDocumentDelivery(actor, operationID)
	if err != nil {
		return DocumentDeliveryRecord{}, err
	}
	if current.MarkdownArtifactID != "" {
		return current, nil
	}
	artifactID := stableDeliveryArtifactID(actor, operationID, artifact.Role)
	ref, err := s.blobs.Put(newID("blob"), artifact.Content)
	if err != nil {
		return current, err
	}
	if fault := s.takeDocumentDeliveryFault(DocumentDeliveryFaultAfterMarkdownBlob); fault != nil {
		_ = s.blobs.Delete(ref)
		return current, fault
	}
	used := false
	err = s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		stored, found, err := s.documentDeliveryForWrite(tx, actor, operationID)
		if err != nil || !found {
			if !found && err == nil {
				err = ErrNotFound
			}
			return err
		}
		if stored.MarkdownArtifactID != "" {
			return nil
		}
		if fault := s.takeDocumentDeliveryFault(DocumentDeliveryFaultBeforeMarkdownCommit); fault != nil {
			return fault
		}
		now := time.Now().UTC()
		if err := insertDeliveryArtifactTx(s, tx, actor, stored.ProjectID, stored.SourceRun, artifactID, ref, artifact, now); err != nil {
			return err
		}
		_, err = tx.Exec(s.q(`UPDATE document_render_deliveries SET state=?,canonical_hash=?,markdown_artifact_id=?,markdown_sha256=?,error_text='',updated_at=? WHERE owner=? AND operation_id=?`),
			DocumentDeliverySourceReady, canonicalHash, artifactID, artifact.SHA256, now.Format(time.RFC3339Nano), actor, operationID)
		if err == nil {
			used = true
		}
		return err
	})
	if !used {
		_ = s.blobs.Delete(ref)
	}
	record, readErr := s.GetDocumentDelivery(actor, operationID)
	if err != nil {
		if used && readErr == nil {
			return record, ErrDocumentDeliveryOutcomeUnknown
		}
		return current, err
	}
	return record, readErr
}

func (s *Store) CommitDocumentDeliveryDerivedAccountChecked(actor, operationID, deliveryKey, profileHash, reviewHash string, artifacts []DocumentDeliveryArtifactInput, allowImplicitDefault bool) (DocumentDeliveryRecord, error) {
	if deliveryKey == "" || profileHash == "" || reviewHash == "" {
		return DocumentDeliveryRecord{}, errors.New("document delivery key, profile hash, and review hash are required")
	}
	ordered, err := validateDerivedArtifacts(artifacts)
	if err != nil {
		return DocumentDeliveryRecord{}, err
	}
	current, err := s.GetDocumentDelivery(actor, operationID)
	if err != nil {
		return DocumentDeliveryRecord{}, err
	}
	if current.State == DocumentDeliveryComplete {
		if current.DeliveryKey != deliveryKey || current.ProfileHash != profileHash || current.ReviewHash != reviewHash {
			return current, ErrDocumentDeliveryConflict
		}
		return current, nil
	}
	if current.MarkdownArtifactID == "" {
		return current, errors.New("document delivery has no authoritative Markdown")
	}
	type stagedArtifact struct {
		input   DocumentDeliveryArtifactInput
		id, ref string
	}
	staged := make([]stagedArtifact, 0, len(ordered))
	cleanup := func() {
		for _, item := range staged {
			_ = s.blobs.Delete(item.ref)
		}
	}
	for _, input := range ordered {
		ref, err := s.blobs.Put(newID("blob"), input.Content)
		if err != nil {
			cleanup()
			return current, err
		}
		staged = append(staged, stagedArtifact{input: input, id: stableDeliveryArtifactID(actor, operationID, input.Role), ref: ref})
		stage := map[DocumentDeliveryArtifactRole]string{DocumentArtifactDOCX: DocumentDeliveryFaultAfterDOCXBlob, DocumentArtifactManifest: DocumentDeliveryFaultAfterManifestBlob, DocumentArtifactArchive: DocumentDeliveryFaultAfterArchiveBlob}[input.Role]
		if fault := s.takeDocumentDeliveryFault(stage); fault != nil {
			cleanup()
			return current, fault
		}
	}
	used := false
	err = s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		stored, found, err := s.documentDeliveryForWrite(tx, actor, operationID)
		if err != nil || !found {
			if !found && err == nil {
				err = ErrNotFound
			}
			return err
		}
		if stored.State == DocumentDeliveryComplete {
			if stored.DeliveryKey != deliveryKey || stored.ProfileHash != profileHash || stored.ReviewHash != reviewHash {
				return ErrDocumentDeliveryConflict
			}
			return nil
		}
		if stored.MarkdownArtifactID == "" {
			return errors.New("document delivery has no authoritative Markdown")
		}
		if fault := s.takeDocumentDeliveryFault(DocumentDeliveryFaultBeforeDerivedCommit); fault != nil {
			return fault
		}
		now := time.Now().UTC()
		for _, item := range staged {
			if err := insertDeliveryArtifactTx(s, tx, actor, stored.ProjectID, stored.SourceRun, item.id, item.ref, item.input, now); err != nil {
				return err
			}
		}
		_, err = tx.Exec(s.q(`UPDATE document_render_deliveries SET state=?,profile_hash=?,review_hash=?,delivery_key=?,docx_artifact_id=?,docx_sha256=?,manifest_artifact_id=?,manifest_sha256=?,archive_artifact_id=?,archive_sha256=?,error_text='',updated_at=? WHERE owner=? AND operation_id=?`),
			DocumentDeliveryComplete, profileHash, reviewHash, deliveryKey,
			staged[0].id, staged[0].input.SHA256, staged[1].id, staged[1].input.SHA256, staged[2].id, staged[2].input.SHA256,
			now.Format(time.RFC3339Nano), actor, operationID)
		if err == nil {
			used = true
		}
		return err
	})
	if !used {
		cleanup()
	}
	record, readErr := s.GetDocumentDelivery(actor, operationID)
	if err != nil {
		if used && readErr == nil {
			return record, ErrDocumentDeliveryOutcomeUnknown
		}
		return current, err
	}
	if fault := s.takeDocumentDeliveryFault(DocumentDeliveryFaultAfterDerivedCommit); fault != nil {
		return record, fmt.Errorf("%w: %v", ErrDocumentDeliveryOutcomeUnknown, fault)
	}
	return record, readErr
}

func (s *Store) MarkDocumentDeliveryFailedAccountChecked(actor, operationID, message string, allowImplicitDefault bool) (DocumentDeliveryRecord, error) {
	err := s.withEffectiveAccountWrite(actor, allowImplicitDefault, func(tx *sql.Tx, _ Account) error {
		stored, found, err := s.documentDeliveryForWrite(tx, actor, operationID)
		if err != nil || !found {
			if !found && err == nil {
				return ErrNotFound
			}
			return err
		}
		if stored.State == DocumentDeliveryComplete {
			return nil
		}
		_, err = tx.Exec(s.q(`UPDATE document_render_deliveries SET state=?,error_text=?,updated_at=? WHERE owner=? AND operation_id=?`),
			DocumentDeliveryFailed, message, nowISO(), actor, operationID)
		return err
	})
	record, readErr := s.GetDocumentDelivery(actor, operationID)
	if err != nil {
		return record, err
	}
	return record, readErr
}

func (s *Store) GetDocumentDelivery(owner, operationID string) (DocumentDeliveryRecord, error) {
	stored, err := scanDocumentDelivery(s.db.QueryRow(s.q(`SELECT `+documentDeliveryColumns+` FROM document_render_deliveries WHERE owner=? AND operation_id=?`), owner, operationID).Scan)
	if errors.Is(err, sql.ErrNoRows) {
		return DocumentDeliveryRecord{}, ErrNotFound
	}
	if err != nil {
		return DocumentDeliveryRecord{}, err
	}
	return s.hydrateDocumentDelivery(stored)
}

func (s *Store) DocumentDeliveryMarkdown(owner, operationID string) ([]byte, error) {
	record, err := s.GetDocumentDelivery(owner, operationID)
	if err != nil {
		return nil, err
	}
	if record.MarkdownArtifactID == "" {
		return nil, ErrNotFound
	}
	return s.ArtifactContent(owner, record.MarkdownArtifactID)
}

func (s *Store) SetDocumentDeliveryFault(stage string, err error) {
	s.documentDeliveryMu.Lock()
	s.documentDeliveryFaultStage, s.documentDeliveryFault = stage, err
	s.documentDeliveryMu.Unlock()
}

func (s *Store) takeDocumentDeliveryFault(stage string) error {
	if stage == "" {
		return nil
	}
	s.documentDeliveryMu.Lock()
	defer s.documentDeliveryMu.Unlock()
	if s.documentDeliveryFaultStage != stage {
		return nil
	}
	err := s.documentDeliveryFault
	s.documentDeliveryFaultStage, s.documentDeliveryFault = "", nil
	return err
}

func (s *Store) documentDeliveryForWrite(tx *sql.Tx, owner, operationID string) (storedDocumentDelivery, bool, error) {
	query := s.q(`SELECT ` + documentDeliveryColumns + ` FROM document_render_deliveries WHERE owner=? AND operation_id=?`)
	if s.Dialect() == "postgres" {
		query += ` FOR UPDATE`
	}
	stored, err := scanDocumentDelivery(tx.QueryRow(query, owner, operationID).Scan)
	if errors.Is(err, sql.ErrNoRows) {
		return storedDocumentDelivery{}, false, nil
	}
	return stored, err == nil, err
}

func scanDocumentDelivery(scan func(...any) error) (storedDocumentDelivery, error) {
	var stored storedDocumentDelivery
	var frozen string
	err := scan(&stored.OperationID, &stored.Owner, &stored.ProjectID, &stored.ConversationID, &stored.SourceRun,
		&stored.RequestSHA256, &stored.State, &frozen, &stored.RendererIdentity, &stored.Mode, &stored.DocumentID,
		&stored.CanonicalHash, &stored.ProfileHash, &stored.ReviewHash, &stored.DeliveryKey,
		&stored.MarkdownArtifactID, &stored.markdownSHA, &stored.docxID, &stored.docxSHA,
		&stored.manifestID, &stored.manifestSHA, &stored.archiveID, &stored.archiveSHA, &stored.Error)
	if err == nil {
		stored.FrozenAt, err = time.Parse(time.RFC3339Nano, frozen)
	}
	return stored, err
}

func (s *Store) hydrateDocumentDelivery(stored storedDocumentDelivery) (DocumentDeliveryRecord, error) {
	record := stored.DocumentDeliveryRecord
	for _, item := range []struct {
		role     DocumentDeliveryArtifactRole
		id, hash string
	}{
		{DocumentArtifactMarkdown, stored.MarkdownArtifactID, stored.markdownSHA},
		{DocumentArtifactDOCX, stored.docxID, stored.docxSHA},
		{DocumentArtifactManifest, stored.manifestID, stored.manifestSHA},
		{DocumentArtifactArchive, stored.archiveID, stored.archiveSHA},
	} {
		if item.id == "" {
			continue
		}
		artifact, err := s.GetArtifact(record.Owner, item.id)
		if err != nil {
			return DocumentDeliveryRecord{}, err
		}
		record.Artifacts = append(record.Artifacts, DocumentDeliveryArtifact{Role: item.role, Artifact: artifact, SHA256: item.hash})
	}
	return record, nil
}

func validateDocumentDeliveryBegin(actor string, input DocumentDeliveryBeginInput) error {
	if actor == "" || !deliveryOperationPattern.MatchString(input.OperationID) || !validDeliveryHash(input.RequestSHA256) ||
		input.ProjectID == "" || input.ConversationID == "" || input.DocumentID == "" || input.RendererIdentity == "" ||
		(input.Mode != "working" && input.Mode != "preview") {
		return errors.New("invalid document delivery begin input")
	}
	return nil
}

func sameDocumentDeliveryInput(record DocumentDeliveryRecord, input DocumentDeliveryBeginInput) bool {
	return record.ProjectID == input.ProjectID && record.ConversationID == input.ConversationID &&
		record.RequestSHA256 == input.RequestSHA256 && record.RendererIdentity == input.RendererIdentity && record.Mode == input.Mode && record.DocumentID == input.DocumentID
}

func validateDeliveryArtifactInput(input DocumentDeliveryArtifactInput) error {
	if input.Name == "" || input.Kind == "" || input.MIME == "" || !validDeliveryHash(input.SHA256) {
		return errors.New("invalid document delivery artifact")
	}
	sum := sha256.Sum256(input.Content)
	if hex.EncodeToString(sum[:]) != input.SHA256 {
		return errors.New("document delivery artifact hash mismatch")
	}
	return nil
}

func validateDerivedArtifacts(inputs []DocumentDeliveryArtifactInput) ([]DocumentDeliveryArtifactInput, error) {
	want := map[DocumentDeliveryArtifactRole]bool{DocumentArtifactDOCX: true, DocumentArtifactManifest: true, DocumentArtifactArchive: true}
	byRole := make(map[DocumentDeliveryArtifactRole]DocumentDeliveryArtifactInput, len(inputs))
	for _, input := range inputs {
		if !want[input.Role] || byRole[input.Role].Role != "" {
			return nil, errors.New("derived delivery requires unique DOCX, manifest, and archive artifacts")
		}
		if err := validateDeliveryArtifactInput(input); err != nil {
			return nil, err
		}
		byRole[input.Role] = input
	}
	if len(byRole) != len(want) {
		return nil, errors.New("derived delivery requires DOCX, manifest, and archive artifacts")
	}
	return []DocumentDeliveryArtifactInput{byRole[DocumentArtifactDOCX], byRole[DocumentArtifactManifest], byRole[DocumentArtifactArchive]}, nil
}

func insertDeliveryArtifactTx(s *Store, tx *sql.Tx, owner, projectID, sourceRun, id, ref string, input DocumentDeliveryArtifactInput, now time.Time) error {
	_, err := tx.Exec(s.q(`INSERT INTO artifacts(id,project_id,owner,name,kind,mime,size,source_run,content_path,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)`),
		id, projectID, owner, input.Name, input.Kind, input.MIME, len(input.Content), sourceRun, ref,
		now.Format(time.RFC3339Nano), now.Format(time.RFC3339Nano))
	return err
}

func stableDeliveryArtifactID(owner, operationID string, role DocumentDeliveryArtifactRole) string {
	sum := sha256.Sum256([]byte(owner + "\x00" + operationID + "\x00" + string(role)))
	return "art_doc_" + hex.EncodeToString(sum[:12])
}

func validDeliveryHash(value string) bool {
	if len(value) != 64 || strings.ToLower(value) != value {
		return false
	}
	_, err := hex.DecodeString(value)
	return err == nil
}