package knowledgeqa

import (
	"context"
	"fmt"
	"strings"
	"time"

	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

const (
	// The MaaS transport admits at most 8 MiB of response bytes. Leave room for
	// every invalid byte to expand to the three-byte UTF-8 replacement rune.
	maxPersistedDebugModelOutputBytes = 32 << 20

	// These match the MeshX accounting contract. The encoded and count caps
	// leave room for usage, TiDB native-JSON metadata, and fixed attempt columns
	// below the default 6-MiB transaction-entry limit.
	maxPersistedAccountingUsageJSONBytes   = (64 << 10) - 1
	maxPersistedSettlementReferences       = 4096
	maxPersistedSettlementReferenceBytes   = 255
	maxPersistedSettlementLineageJSONBytes = 4 << 20
)

func prepareCompletion(completion Completion) (Completion, []byte, []byte, error) {
	completion.Now = normalizeNow(completion.Now)
	if !isTerminalTurnStatus(completion.Status) {
		return Completion{}, nil, nil, ErrInvalidState
	}
	if completion.RuntimeContext != nil {
		completion.RuntimeContext.Summary = strings.TrimSpace(completion.RuntimeContext.Summary)
		completion.RuntimeContext.CompactedThroughTurnID = strings.TrimSpace(completion.RuntimeContext.CompactedThroughTurnID)
		completion.RuntimeContext.SummaryVersion = strings.TrimSpace(completion.RuntimeContext.SummaryVersion)
		if completion.RuntimeContext.ExpectedGeneration < 0 ||
			completion.RuntimeContext.Summary == "" || len(completion.RuntimeContext.Summary) > 1<<20 ||
			completion.RuntimeContext.CompactedThroughTurnID == "" ||
			completion.RuntimeContext.SummaryVersion == "" || len(completion.RuntimeContext.SummaryVersion) > 64 {
			return Completion{}, nil, nil, fmt.Errorf("invalid Knowledge Q&A runtime context update")
		}
	}
	if completion.Status != "completed" {
		completion.AssistantMessageID = ""
	}
	if len(completion.ErrorDetail) > 2048 {
		completion.ErrorDetail = completion.ErrorDetail[:2048]
	}
	if completion.Status != "failed" {
		completion.DebugModelOutput = ""
	} else {
		completion.DebugModelOutput = strings.ToValidUTF8(completion.DebugModelOutput, "�")
		if len(completion.DebugModelOutput) > maxPersistedDebugModelOutputBytes {
			return Completion{}, nil, nil, fmt.Errorf("Knowledge Q&A debug model output exceeds the runtime response boundary")
		}
	}
	completion.SettlementRefs = normalizeStringSet(completion.SettlementRefs)
	completion.SettlementRef = strings.TrimSpace(completion.SettlementRef)
	if completion.SettlementRef != "" {
		found := false
		for _, settlementRef := range completion.SettlementRefs {
			if settlementRef == completion.SettlementRef {
				found = true
				break
			}
		}
		if !found {
			completion.SettlementRefs = append(completion.SettlementRefs, completion.SettlementRef)
		}
	} else if len(completion.SettlementRefs) > 0 {
		completion.SettlementRef = completion.SettlementRefs[len(completion.SettlementRefs)-1]
	}
	if err := validatePersistedSettlementReferences(completion.SettlementRefs); err != nil {
		return Completion{}, nil, nil, err
	}
	usageJSON, err := model.MarshalJSONObject("Knowledge Q&A usage", nonNilMap(completion.Usage))
	if err != nil {
		return Completion{}, nil, nil, fmt.Errorf("marshal Knowledge Q&A usage: %w", err)
	}
	settlementRefsJSON, err := model.MarshalJSONObject("Knowledge Q&A settlement references", completion.SettlementRefs)
	if err != nil {
		return Completion{}, nil, nil, fmt.Errorf("marshal Knowledge Q&A settlement references: %w", err)
	}
	if err := validatePersistedAccountingDocumentSizes(
		usageJSON, settlementRefsJSON, len(completion.SettlementRefs),
	); err != nil {
		return Completion{}, nil, nil, err
	}
	return completion, usageJSON, settlementRefsJSON, nil
}

func validatePersistedSettlementReferences(settlementRefs []string) error {
	for _, settlementRef := range settlementRefs {
		if len(settlementRef) > maxPersistedSettlementReferenceBytes {
			return fmt.Errorf("Knowledge Q&A settlement lineage exceeds its storage boundary")
		}
	}
	return nil
}

func validatePersistedAccountingDocumentSizes(
	usageJSON, settlementRefsJSON []byte,
	settlementRefCount int,
) error {
	if len(usageJSON) > maxPersistedAccountingUsageJSONBytes {
		return fmt.Errorf("Knowledge Q&A accounting usage exceeds its storage boundary")
	}
	if settlementRefCount > maxPersistedSettlementReferences ||
		len(settlementRefsJSON) > maxPersistedSettlementLineageJSONBytes {
		return fmt.Errorf("Knowledge Q&A settlement lineage exceeds its storage boundary")
	}
	return nil
}

func applyRuntimeContextUpdateTx(
	ctx context.Context,
	tx coredb.Tx,
	workspaceID, sessionID, currentTurnID string,
	update RuntimeContextUpdate,
	now time.Time,
) error {
	var compactedTurnCount int
	if err := tx.QueryRow(ctx, `
SELECT COUNT(*)
FROM knowledge_qa_turns AS compacted
JOIN knowledge_qa_turns AS current_turn
  ON current_turn.workspace_id=compacted.workspace_id
  AND current_turn.session_id=compacted.session_id
  AND current_turn.id=$3
WHERE compacted.workspace_id=$1 AND compacted.session_id=$2
  AND compacted.id=$4
  AND compacted.status IN ('completed','no_answer','failed','canceled')
  AND (compacted.queued_at < current_turn.queued_at OR
       (compacted.queued_at=current_turn.queued_at AND compacted.id < current_turn.id))
`, workspaceID, sessionID, currentTurnID, update.CompactedThroughTurnID).Scan(&compactedTurnCount); err != nil {
		return fmt.Errorf("validate Knowledge Q&A runtime context boundary: %w", err)
	}
	if compactedTurnCount != 1 {
		return ErrInvalidState
	}
	var generation int64
	err := tx.QueryRow(ctx, `
SELECT generation FROM knowledge_qa_runtime_contexts
WHERE session_id=$1 AND workspace_id=$2
FOR UPDATE
`, sessionID, workspaceID).Scan(&generation)
	switch {
	case coredb.IsNotFound(err):
		if update.ExpectedGeneration != 0 {
			return ErrInvalidState
		}
		if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_runtime_contexts (
  session_id, workspace_id, generation, summary,
  compacted_through_turn_id, summary_version, created_at, updated_at
) VALUES ($1,$2,1,$3,$4,$5,$6,$6)
`, sessionID, workspaceID, update.Summary, update.CompactedThroughTurnID,
			update.SummaryVersion, now); err != nil {
			return fmt.Errorf("create Knowledge Q&A runtime context: %w", err)
		}
	case err != nil:
		return fmt.Errorf("lock Knowledge Q&A runtime context: %w", err)
	default:
		if generation != update.ExpectedGeneration {
			return ErrInvalidState
		}
		result, err := tx.Exec(ctx, `
UPDATE knowledge_qa_runtime_contexts
SET generation=generation+1, summary=$3, compacted_through_turn_id=$4,
    summary_version=$5, updated_at=$6
WHERE session_id=$1 AND workspace_id=$2 AND generation=$7
`, sessionID, workspaceID, update.Summary, update.CompactedThroughTurnID,
			update.SummaryVersion, now, update.ExpectedGeneration)
		if err != nil {
			return fmt.Errorf("update Knowledge Q&A runtime context: %w", err)
		}
		if result.RowsAffected() != 1 {
			return ErrInvalidState
		}
	}
	return nil
}