Agent-System/internal/docgov/lease.go

package docgov

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

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

// 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
	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) {
	e.mu.Lock()
	defer e.mu.Unlock()

	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")
	}

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

	now := time.Now()
	ttl := req.TTL
	if ttl <= 0 {
		if req.HolderType == LeaseHolderSquad {
			ttl = DefaultSquadLeaseTTL
		} else {
			ttl = DefaultHumanLeaseTTL
		}
	}

	if doc.HasActiveLease(now) {
		sameHolder := doc.LeaseHolderType == req.HolderType && doc.LeaseHolderID == req.HolderID
		if !sameHolder && req.HolderType == LeaseHolderSquad && e.shouldReclaimSquadLease(ctx, doc) {
			e.appendAuditLocked(doc, "lease_reclaim_orphan", 0, string(req.HolderType), RoleSystem, map[string]any{
				"previous_holder_type": doc.LeaseHolderType,
				"previous_holder_id":   doc.LeaseHolderID,
			})
			doc.ClearLease()
			e.persistLeaseLocked(ctx, doc)
		} else if !sameHolder {
			e.appendAuditLocked(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,
			})
			return doc, &DocumentBusyError{
				DocID:      doc.DocID,
				HolderType: doc.LeaseHolderType,
				HolderID:   doc.LeaseHolderID,
				WorkDir:    doc.LeaseWorkDir,
				ExpiresAt:  *doc.LeaseExpiresAt,
				Reason:     fmt.Sprintf("文档【%s】正在被 %s(%s) 编辑中，请稍后再试", doc.Title, doc.LeaseHolderType, doc.LeaseHolderID),
			}
		}
	}

	expires := now.Add(ttl)
	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

	e.persistLeaseLocked(ctx, doc)
	e.appendAuditLocked(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) {
	e.mu.Lock()
	defer e.mu.Unlock()

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

	now := time.Now()
	if !doc.HasActiveLease(now) {
		doc.ClearLease()
		doc.UpdatedAt = now
		e.persistLeaseLocked(ctx, doc)
		return doc, nil
	}

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

	doc.ClearLease()
	doc.UpdatedAt = now
	e.persistLeaseLocked(ctx, doc)
	e.appendAuditLocked(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) {
	e.mu.Lock()
	defer e.mu.Unlock()

	doc, err := e.getDocLocked(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
	}

	if ttl <= 0 {
		if doc.LeaseHolderType == LeaseHolderSquad {
			ttl = DefaultSquadLeaseTTL
		} else {
			ttl = DefaultHumanLeaseTTL
		}
	}
	expires := now.Add(ttl)
	doc.LeaseExpiresAt = &expires
	doc.LeaseHeartbeatAt = &now
	doc.UpdatedAt = now
	e.persistLeaseLocked(ctx, doc)
	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 {
	e.mu.Lock()
	defer e.mu.Unlock()

	doc, err := e.getDocLocked(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) {
	e.mu.Lock()
	defer e.mu.Unlock()

	for _, p := range e.memProps {
		if p.DocID == docID && p.Status == ProposalApprovedByAuthor {
			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.memProps[p.ProposalID] = &p
			return &p, nil
		}
	}
	return nil, nil
}

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

	p, ok := e.memProps[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.memProps[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
}

// AssertHumanCanEdit checks author duty + 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) error {
	e.mu.Lock()
	doc, err := e.getDocLocked(ctx, docID)
	if err != nil {
		e.mu.Unlock()
		return err
	}
	if err := e.sm.CanEditDocumentContent(doc, actorID, actorRole); err != nil {
		e.mu.Unlock()
		return err
	}
	e.mu.Unlock()

	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.getDocLocked(ctx, req.DocID)
	if err != nil {
		return nil, err
	}

	if err := e.sm.CanEditDocumentContent(doc, req.ActorID, req.ActorRole); 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, "output.md")
		if werr := fsutil.WriteFileAtomic(outPath, []byte(canon), 0o644); werr != nil {
			return nil, fmt.Errorf("docgov: write content failed: %w", werr)
		}
	}

	if e.pool != nil {
		_, _ = 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)
	}

	e.appendAuditLocked(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 {
	e.mu.Lock()
	defer e.mu.Unlock()
	doc, err := e.getDocLocked(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.memDocs[docID]
	if !exists {
		var err error
		doc, err = e.loadDocumentFromDBLocked(ctx, docID)
		if err != nil || doc == nil {
			return nil, ErrDocumentNotFound
		}
	}
	return doc, nil
}

func (e *GovernanceEngine) persistLeaseLocked(ctx context.Context, doc *DocumentRecord) {
	if e.pool == nil {
		return
	}
	_, _ = 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
	`, string(doc.LeaseHolderType), doc.LeaseHolderID, doc.LeaseWorkDir,
		doc.LeaseExpiresAt, doc.LeaseHeartbeatAt, doc.LeaseFencingToken, doc.UpdatedAt, doc.DocID)
}

func (e *GovernanceEngine) appendAuditLocked(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.memAudits[doc.DocID] = append(e.memAudits[doc.DocID], audit)
	if e.pool != nil {
		detJSON, _ := json.Marshal(details)
		_, _ = e.pool.Exec(context.Background(), `
			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)
	}
}