Agent-System/internal/docgov/lease.go

package docgov

import (
	"context"
	"crypto/rand"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io"
	"log/slog"
	"path/filepath"
	"strings"
	"time"

	"chat-agent-system/internal/doccontent"
	"chat-agent-system/internal/fsutil"
	"chat-agent-system/internal/workspace"
)

// LeaseAcquireRequest 申请编辑租约
type LeaseAcquireRequest struct {
	DocID      string
	HolderType LeaseHolderType
	HolderID   string
	WorkDir    string
	TTL        time.Duration
}

// ContentSaveRequest 正文/块内容保存请求
type ContentSaveRequest struct {
	DocID          string
	ActorID        int64
	ActorRole      UserRole
	ActorName      string
	ActorIsMember  bool // server 层已用 authorizer 校验的项目成员；不得用 RoleSystem 冒充
	MarkdownBody   string
	ExpectedSHA256 string // 乐观并发：非空时必须匹配当前 ContentSHA256
	FencingToken   string
}

// ContentSaveResult 正文保存结果
type ContentSaveResult struct {
	Doc           *DocumentRecord `json:"doc"`
	ContentSHA256 string          `json:"content_sha256"`
	WorkDir       string          `json:"work_dir,omitempty"`
}

// AcquireLease 获取文档编辑租约；同 holder 续租，异 holder 且未过期返回 ErrDocumentBusy
func (e *GovernanceEngine) AcquireLease(ctx context.Context, req LeaseAcquireRequest) (*DocumentRecord, error) {
	if req.HolderType != LeaseHolderSquad && req.HolderType != LeaseHolderHuman {
		return nil, fmt.Errorf("docgov: invalid lease holder type %q", req.HolderType)
	}
	if strings.TrimSpace(req.HolderID) == "" {
		return nil, fmt.Errorf("docgov: lease holder id is required")
	}

	// T20 B3: doc 级锁取代全局锁；DB RTT 只阻塞同文档。
	g := e.acquireDocLock(req.DocID)
	defer g.Release()

	doc, err := e.ensureDocLockedWithDB(ctx, req.DocID)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	ttl := clampLeaseTTL(req.HolderType, req.TTL)

	if doc.HasActiveLease(now) {
		sameHolder := doc.LeaseHolderType == req.HolderType && doc.LeaseHolderID == req.HolderID
		if !sameHolder && req.HolderType == LeaseHolderSquad && e.shouldReclaimSquadLease(ctx, doc) {
			e.appendAuditLocked(ctx, doc, "lease_reclaim_orphan", 0, string(req.HolderType), RoleSystem, map[string]any{
				"previous_holder_type": doc.LeaseHolderType,
				"previous_holder_id":   doc.LeaseHolderID,
			})
			doc.ClearLease()
			if perr := e.persistLeaseLocked(ctx, doc); perr != nil {
				slog.Error("docgov_lease_reclaim_persist_failed", "doc_id", doc.DocID, "error", perr)
			}
		} else if !sameHolder {
			e.appendAuditLocked(ctx, doc, "lease_deny", 0, string(req.HolderType), RoleSystem, map[string]any{
				"requester_type": req.HolderType,
				"requester_id":   req.HolderID,
				"holder_type":    doc.LeaseHolderType,
				"holder_id":      doc.LeaseHolderID,
				"expires_at":     doc.LeaseExpiresAt,
			})
			display := LeaseHolderDisplay(doc.LeaseHolderType, doc.LeaseHolderID, "")
			return doc, &DocumentBusyError{
				DocID:             doc.DocID,
				HolderType:        doc.LeaseHolderType,
				HolderID:          doc.LeaseHolderID,
				HolderDisplayName: display,
				WorkDir:           doc.LeaseWorkDir,
				ExpiresAt:         *doc.LeaseExpiresAt,
				Reason:            FormatDocumentBusyReason(doc.Title, display, *doc.LeaseExpiresAt),
			}
		}
	}

	expires := now.Add(ttl)
	// 先持久化再更新内存，失败时不污染内存租约状态
	prevType, prevID, prevWD := doc.LeaseHolderType, doc.LeaseHolderID, doc.LeaseWorkDir
	prevExpires, prevHeartbeat, prevFencing := doc.LeaseExpiresAt, doc.LeaseHeartbeatAt, doc.LeaseFencingToken
	doc.LeaseHolderType = req.HolderType
	doc.LeaseHolderID = req.HolderID
	if req.WorkDir != "" {
		doc.LeaseWorkDir = req.WorkDir
	}
	doc.LeaseExpiresAt = &expires
	doc.LeaseHeartbeatAt = &now
	doc.LeaseFencingToken = newLeaseFencingToken()
	doc.UpdatedAt = now

	if perr := e.persistLeaseLocked(ctx, doc); perr != nil {
		doc.LeaseHolderType, doc.LeaseHolderID, doc.LeaseWorkDir = prevType, prevID, prevWD
		doc.LeaseExpiresAt, doc.LeaseHeartbeatAt, doc.LeaseFencingToken = prevExpires, prevHeartbeat, prevFencing
		return nil, perr
	}
	e.appendAuditLocked(ctx, doc, "lease_acquire", 0, string(req.HolderType), RoleSystem, map[string]any{
		"holder_type": req.HolderType,
		"holder_id":   req.HolderID,
		"work_dir":    doc.LeaseWorkDir,
		"expires_at":  expires,
		"ttl":         ttl.String(),
	})

	return doc, nil
}

// ReleaseLease 释放编辑租约（仅持有者或过期后任意调用清理）
func (e *GovernanceEngine) ReleaseLease(ctx context.Context, docID, holderType, holderID string) (*DocumentRecord, error) {
	g := e.acquireDocLock(docID)
	defer g.Release()

	doc, err := e.ensureDocLockedWithDB(ctx, docID)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	if !doc.HasActiveLease(now) {
		doc.ClearLease()
		doc.UpdatedAt = now
		if perr := e.persistLeaseLocked(ctx, doc); perr != nil {
			slog.Error("docgov_lease_release_persist_failed", "doc_id", doc.DocID, "error", perr)
		}
		return doc, nil
	}

	if LeaseHolderType(holderType) != doc.LeaseHolderType || holderID != doc.LeaseHolderID {
		return doc, ErrLeaseNotHeld
	}

	doc.ClearLease()
	doc.UpdatedAt = now
	if perr := e.persistLeaseLocked(ctx, doc); perr != nil {
		slog.Error("docgov_lease_release_persist_failed", "doc_id", doc.DocID, "error", perr)
	}
	e.appendAuditLocked(ctx, doc, "lease_release", 0, holderType, RoleSystem, map[string]any{
		"holder_type": holderType,
		"holder_id":   holderID,
	})
	return doc, nil
}

// HeartbeatLease 续租心跳
func (e *GovernanceEngine) HeartbeatLease(ctx context.Context, docID, holderType, holderID string, ttl time.Duration) (*DocumentRecord, error) {
	g := e.acquireDocLock(docID)
	defer g.Release()

	doc, err := e.ensureDocLockedWithDB(ctx, docID)
	if err != nil {
		return nil, err
	}

	now := time.Now()
	if !doc.HasActiveLease(now) {
		return doc, ErrLeaseExpired
	}
	if LeaseHolderType(holderType) != doc.LeaseHolderType || holderID != doc.LeaseHolderID {
		return doc, ErrLeaseNotHeld
	}

	ttl = clampLeaseTTL(doc.LeaseHolderType, ttl)
	expires := now.Add(ttl)
	doc.LeaseExpiresAt = &expires
	doc.LeaseHeartbeatAt = &now
	doc.UpdatedAt = now
	if perr := e.persistLeaseLocked(ctx, doc); perr != nil {
		slog.Error("docgov_lease_heartbeat_persist_failed", "doc_id", doc.DocID, "error", perr)
	}
	return doc, nil
}

// shouldReclaimSquadLease 当 squad 租约 holder 无对应 running task_run 时视为孤儿锁并回收。
func (e *GovernanceEngine) shouldReclaimSquadLease(ctx context.Context, doc *DocumentRecord) bool {
	if doc == nil || doc.LeaseHolderType != LeaseHolderSquad {
		return false
	}
	holderID := strings.TrimSpace(doc.LeaseHolderID)
	if holderID == "" || !strings.HasPrefix(holderID, "run_") {
		return false
	}
	if e.pool == nil {
		return false
	}
	var status string
	err := e.pool.QueryRow(ctx, `SELECT status FROM task_runs WHERE run_id = $1`, holderID).Scan(&status)
	if err != nil {
		return true
	}
	return status != "running"
}

// AssertWritableForComment 批注前检查：有效 edit lease 期间禁止批注
func (e *GovernanceEngine) AssertWritableForComment(ctx context.Context, docID string) error {
	g := e.acquireDocLock(docID)
	defer g.Release()

	doc, err := e.ensureDocLockedWithDB(ctx, docID)
	if err != nil {
		return err
	}
	now := time.Now()
	if doc.HasActiveLease(now) {
		return &DocumentBusyError{
			DocID:      doc.DocID,
			HolderType: doc.LeaseHolderType,
			HolderID:   doc.LeaseHolderID,
			WorkDir:    doc.LeaseWorkDir,
			ExpiresAt:  *doc.LeaseExpiresAt,
			Reason:     fmt.Sprintf("文档正在由 %s 编辑/生成中，暂不可批注", doc.LeaseHolderType),
		}
	}
	return nil
}

// FindApprovedProposal 查找编制人已同意、尚未 applied 的提案
func (e *GovernanceEngine) FindApprovedProposal(ctx context.Context, docID string) (*ChangeProposal, error) {
	g := e.acquireDocLock(docID)
	defer g.Release()

	p := e.memPropsFind(func(p *ChangeProposal) bool {
		return p.DocID == docID && p.Status == ProposalApprovedByAuthor
	})
	if p != nil {
		return p, nil
	}
	if e.pool != nil {
		var p ChangeProposal
		var stat string
		err := e.pool.QueryRow(ctx, `
			SELECT id, proposal_id, doc_id, project_id, triggered_by, proposed_changes, impact_analysis, risk_assessment, status, author_decision_reason, decided_at, created_at
			FROM document_change_proposals
			WHERE doc_id = $1 AND status = 'approved_by_author'
			ORDER BY decided_at DESC NULLS LAST, created_at DESC
			LIMIT 1
		`, docID).Scan(&p.ID, &p.ProposalID, &p.DocID, &p.ProjectID, &p.TriggeredBy, &p.ProposedChanges, &p.ImpactAnalysis, &p.RiskAssessment, &stat, &p.AuthorDecisionReason, &p.DecidedAt, &p.CreatedAt)
		if err == nil {
			p.Status = ProposalStatus(stat)
			e.memPropsPut(p.ProposalID, &p)
			return &p, nil
		}
	}
	return nil, nil
}

// MarkProposalApplied 将提案标记为已应用。
// DB 写成功后再更新内存，避免 DB 失败时内存已假 applied；已 applied 则幂等成功。
func (e *GovernanceEngine) MarkProposalApplied(ctx context.Context, proposalID string) error {
	g := e.acquireDocLock("*proposal:" + proposalID)
	defer g.Release()

	p, ok := e.memPropsGet(proposalID)
	if !ok && e.pool != nil {
		var loaded ChangeProposal
		var stat string
		err := e.pool.QueryRow(ctx, `
			SELECT id, proposal_id, doc_id, project_id, triggered_by, proposed_changes, impact_analysis, risk_assessment, status, author_decision_reason, decided_at, created_at
			FROM document_change_proposals WHERE proposal_id = $1
		`, proposalID).Scan(&loaded.ID, &loaded.ProposalID, &loaded.DocID, &loaded.ProjectID, &loaded.TriggeredBy, &loaded.ProposedChanges, &loaded.ImpactAnalysis, &loaded.RiskAssessment, &stat, &loaded.AuthorDecisionReason, &loaded.DecidedAt, &loaded.CreatedAt)
		if err != nil {
			return ErrProposalNotFound
		}
		loaded.Status = ProposalStatus(stat)
		p = &loaded
		e.memPropsPut(proposalID, p)
	}
	if p == nil {
		return ErrProposalNotFound
	}
	if p.Status == ProposalApplied {
		return nil // idempotent
	}
	if p.Status != ProposalApprovedByAuthor {
		return fmt.Errorf("docgov: proposal %s is not approved (status=%s)", proposalID, p.Status)
	}

	if e.pool != nil {
		tag, err := e.pool.Exec(ctx, `
			UPDATE document_change_proposals SET status = $1 WHERE proposal_id = $2 AND status = $3
		`, string(ProposalApplied), proposalID, string(ProposalApprovedByAuthor))
		if err != nil {
			return fmt.Errorf("docgov: mark proposal applied: %w", err)
		}
		if tag.RowsAffected() == 0 {
			// Concurrent apply or already applied in DB — reload and reconcile memory.
			var stat string
			qErr := e.pool.QueryRow(ctx, `SELECT status FROM document_change_proposals WHERE proposal_id = $1`, proposalID).Scan(&stat)
			if qErr != nil {
				return fmt.Errorf("docgov: mark proposal applied: no rows updated for %s", proposalID)
			}
			if ProposalStatus(stat) == ProposalApplied {
				p.Status = ProposalApplied
				return nil
			}
			return fmt.Errorf("docgov: mark proposal applied: no rows updated for %s (status=%s)", proposalID, stat)
		}
	}

	// Memory only after durable write succeeds (or in-memory engine with no pool).
	p.Status = ProposalApplied
	return nil
}

func clampLeaseTTL(holder LeaseHolderType, ttl time.Duration) time.Duration {
	maxTTL := DefaultHumanLeaseTTL
	if holder == LeaseHolderSquad {
		maxTTL = DefaultSquadLeaseTTL
	}
	if ttl <= 0 || ttl > maxTTL {
		return maxTTL
	}
	return ttl
}

// CheckHumanCanEdit verifies drafting state and edit permission without taking a lease.
// actorIsMember must be set by the server after project-membership / DocAccessActionEdit checks.
func (e *GovernanceEngine) CheckHumanCanEdit(ctx context.Context, docID string, actorID int64, actorRole UserRole, actorIsMember bool) error {
	g := e.acquireDocLock(docID)
	defer g.Release()
	doc, err := e.ensureDocLockedWithDB(ctx, docID)
	if err != nil {
		return err
	}
	return e.sm.CanEditDocumentContent(doc, actorID, actorRole, actorIsMember)
}

// AssertHumanCanEdit checks edit permission + drafting lifecycle BEFORE acquiring/renewing
// the human lease, so unauthorized actors never become lease holders.
func (e *GovernanceEngine) AssertHumanCanEdit(ctx context.Context, docID string, actorID int64, actorRole UserRole, actorIsMember bool) error {
	if err := e.CheckHumanCanEdit(ctx, docID, actorID, actorRole, actorIsMember); err != nil {
		return err
	}

	holderID := fmt.Sprintf("%d", actorID)
	if actorID <= 0 {
		holderID = "human_anon"
	}
	_, err := e.AcquireLease(ctx, LeaseAcquireRequest{
		DocID:      docID,
		HolderType: LeaseHolderHuman,
		HolderID:   holderID,
		TTL:        DefaultHumanLeaseTTL,
	})
	return err
}

// SaveDocumentContent 人工正文保存：校验可编辑 + 获取/续 human 租约 + 乐观哈希 + 落盘
func (e *GovernanceEngine) SaveDocumentContent(ctx context.Context, req ContentSaveRequest) (*ContentSaveResult, error) {
	holderID := fmt.Sprintf("%d", req.ActorID)
	if req.ActorID <= 0 {
		holderID = "human_anon"
	}
	if req.FencingToken != "" {
		if err := e.AssertLeaseFencing(ctx, req.DocID, req.FencingToken); err != nil {
			return nil, err
		}
	}

	doc, err := e.AcquireLease(ctx, LeaseAcquireRequest{
		DocID:      req.DocID,
		HolderType: LeaseHolderHuman,
		HolderID:   holderID,
		TTL:        DefaultHumanLeaseTTL,
	})
	if err != nil {
		return nil, err
	}

	e.mu.Lock()
	defer e.mu.Unlock()

	doc, err = e.ensureDocLockedWithDB(ctx, req.DocID)
	if err != nil {
		return nil, err
	}

	if err := e.sm.CanEditDocumentContent(doc, req.ActorID, req.ActorRole, req.ActorIsMember); err != nil {
		return nil, err
	}

	if req.ExpectedSHA256 != "" && doc.ContentSHA256 != "" && req.ExpectedSHA256 != doc.ContentSHA256 {
		return nil, fmt.Errorf("docgov: content version conflict: expected sha %s, current %s", req.ExpectedSHA256, doc.ContentSHA256)
	}

	now := time.Now()
	canon := doccontent.CanonicalizeMarkdown(req.MarkdownBody)
	shaStr := doccontent.SHA256Hex(canon)
	doc.ContentSHA256 = shaStr
	doc.IsTemplateOnly = false
	doc.AuthorClaimed = true
	doc.AuthorLastEditedAt = &now
	if doc.AuthorClaimedAt == nil {
		doc.AuthorClaimedAt = &now
	}
	doc.UpdatedAt = now

	workDir := doc.WorkDir
	if workDir != "" {
		outPath := filepath.Join(workDir, workspace.DeliverableKey)
		if werr := fsutil.WriteFileAtomic(outPath, []byte(canon), 0o644); werr != nil {
			return nil, fmt.Errorf("docgov: write content failed: %w", werr)
		}
	}

	if e.pool != nil {
		if _, perr := e.pool.Exec(ctx, `
			UPDATE document_records
			SET content_sha256 = $1, is_template_only = $2, work_dir = $3,
			    author_claimed = $4, author_claimed_at = $5, author_last_edited_at = $6, updated_at = $7
			WHERE doc_id = $8
		`, doc.ContentSHA256, doc.IsTemplateOnly, doc.WorkDir,
			doc.AuthorClaimed, doc.AuthorClaimedAt, doc.AuthorLastEditedAt, doc.UpdatedAt, doc.DocID); perr != nil {
			slog.Error("docgov_content_persist_failed", "doc_id", doc.DocID, "error", perr)
			return nil, fmt.Errorf("docgov: persist content: %w", perr)
		}
	}

	e.appendAuditLocked(ctx, doc, "author_edit_content", req.ActorID, req.ActorName, req.ActorRole, map[string]any{
		"content_sha256": shaStr,
		"bytes":          len(canon),
	})

	return &ContentSaveResult{Doc: doc, ContentSHA256: shaStr, WorkDir: workDir}, nil
}

// AssertLeaseFencing rejects writers holding a superseded token.
func (e *GovernanceEngine) AssertLeaseFencing(ctx context.Context, docID, token string) error {
	g := e.acquireDocLock(docID)
	defer g.Release()
	doc, err := e.ensureDocLockedWithDB(ctx, docID)
	if err != nil {
		return err
	}
	if token == "" || doc.LeaseFencingToken == "" || token != doc.LeaseFencingToken {
		return ErrStaleFencing
	}
	return nil
}

func newLeaseFencingToken() string {
	var b [16]byte
	if _, err := io.ReadFull(rand.Reader, b[:]); err != nil {
		return fmt.Sprintf("%d", time.Now().UnixNano())
	}
	return hex.EncodeToString(b[:])
}

func (e *GovernanceEngine) getDocLocked(ctx context.Context, docID string) (*DocumentRecord, error) {
	doc, exists := e.memDocsGet(docID)
	if !exists {
		var err error
		doc, err = e.loadDocumentFromDBLocked(ctx, docID)
		if err != nil || doc == nil {
			return nil, ErrDocumentNotFound
		}
	}
	return doc, nil
}

// LeaseHolderDisplay returns a user-facing holder label.
// squad → 项目助手; human → display name if provided, else "其他用户".
func LeaseHolderDisplay(holderType LeaseHolderType, holderID, displayName string) string {
	switch holderType {
	case LeaseHolderSquad:
		return "项目助手"
	case LeaseHolderHuman:
		if strings.TrimSpace(displayName) != "" {
			return strings.TrimSpace(displayName)
		}
		return "其他用户"
	default:
		if strings.TrimSpace(displayName) != "" {
			return strings.TrimSpace(displayName)
		}
		return string(holderType)
	}
}

// FormatDocumentBusyReason builds the user-visible busy copy.
func FormatDocumentBusyReason(title, holderDisplay string, expiresAt time.Time) string {
	remain := "稍后"
	if !expiresAt.IsZero() {
		m := int(time.Until(expiresAt).Minutes())
		if m < 1 {
			m = 1
		}
		remain = fmt.Sprintf("%d 分钟后", m)
	}
	t := strings.TrimSpace(title)
	if t == "" {
		t = "文档"
	}
	h := strings.TrimSpace(holderDisplay)
	if h == "" {
		h = "其他任务"
	}
	return fmt.Sprintf("《%s》正在修改中（由%s发起，预计 %s释放）。您的修改请求未执行，请稍后重试。", t, h, remain)
}

func (e *GovernanceEngine) persistLeaseLocked(ctx context.Context, doc *DocumentRecord) error {
	if e.pool == nil {
		return nil
	}
	// Release/clear: write empty lease unconditionally (caller already authorized).
	if doc.LeaseHolderType == "" || doc.LeaseHolderType == LeaseHolderNone {
		_, err := e.pool.Exec(ctx, `
			UPDATE document_records
			SET lease_holder_type = '', lease_holder_id = '', lease_work_dir = $1,
			    lease_expires_at = NULL, lease_heartbeat_at = NULL, lease_fencing_token = '', updated_at = $2
			WHERE doc_id = $3
		`, doc.LeaseWorkDir, doc.UpdatedAt, doc.DocID)
		if err != nil {
			slog.Error("docgov_lease_clear_persist_failed", "doc_id", doc.DocID, "error", err)
			return fmt.Errorf("docgov: persist lease clear: %w", err)
		}
		return nil
	}
	// Acquire/renew: conditional — free, expired, or same holder only.
	tag, err := e.pool.Exec(ctx, `
		UPDATE document_records
		SET lease_holder_type = $1, lease_holder_id = $2, lease_work_dir = $3,
		    lease_expires_at = $4, lease_heartbeat_at = $5, lease_fencing_token = $6, updated_at = $7
		WHERE doc_id = $8
		  AND (
		    lease_holder_type = '' OR lease_holder_type = 'none'
		    OR lease_expires_at IS NULL OR lease_expires_at < now()
		    OR (lease_holder_type = $1 AND lease_holder_id = $2)
		  )
	`, string(doc.LeaseHolderType), doc.LeaseHolderID, doc.LeaseWorkDir,
		doc.LeaseExpiresAt, doc.LeaseHeartbeatAt, doc.LeaseFencingToken, doc.UpdatedAt, doc.DocID)
	if err != nil {
		slog.Error("docgov_lease_persist_failed", "doc_id", doc.DocID, "holder", doc.LeaseHolderID, "error", err)
		return fmt.Errorf("docgov: persist lease: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrDocumentBusy
	}
	return nil
}

func (e *GovernanceEngine) appendAuditLocked(ctx context.Context, doc *DocumentRecord, action string, actorID int64, actorName string, actorRole UserRole, details map[string]any) {
	now := time.Now()
	audit := &AuditLog{
		DocID:     doc.DocID,
		Action:    action,
		ActorID:   actorID,
		ActorName: actorName,
		ActorRole: actorRole,
		FromState: doc.CurrentState,
		ToState:   doc.CurrentState,
		Details:   details,
		CreatedAt: now,
	}
	e.memAuditsAppend(doc.DocID, audit)
	if e.pool != nil {
		detJSON, _ := json.Marshal(details)
		if _, err := e.pool.Exec(ctx, `
			INSERT INTO document_audit_logs (doc_id, action, actor_id, actor_name, actor_role, from_state, to_state, details, created_at)
			VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9)
		`, audit.DocID, audit.Action, audit.ActorID, audit.ActorName, string(audit.ActorRole),
			string(audit.FromState), string(audit.ToState), detJSON, audit.CreatedAt); err != nil {
			slog.Error("docgov_audit_persist_failed", "doc_id", audit.DocID, "action", audit.Action, "error", err)
		}
	}
}