package knowledgeqa

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

	"github.com/tidbcloud/AgentTeam/internal/core/identifiers"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
)

// PrepareCitations normalizes citation metadata and allocates stable IDs
// before the assistant message is rendered. The prepared values can then be
// committed with that message at the turn's terminal serialization point.
func PrepareCitations(citations []Citation) ([]Citation, error) {
	prepared := make([]Citation, 0, len(citations))
	for _, citation := range citations {
		citation = normalizeCitation(citation)
		if citation.ID == "" {
			id, err := identifiers.NewEntityID("kqc")
			if err != nil {
				return nil, fmt.Errorf("generate Knowledge Q&A citation id: %w", err)
			}
			citation.ID = id
		}
		if !validCitation(citation) {
			return nil, fmt.Errorf("invalid Knowledge Q&A citation")
		}
		if len(citation.ResolvedRef) > 255 ||
			(citation.ResolvedRefKind != "" && citation.ResolvedRefKind != "branch" && citation.ResolvedRefKind != "tag") ||
			((citation.ResolvedRef == "") != (citation.ResolvedRefKind == "")) {
			return nil, fmt.Errorf("invalid Knowledge Q&A citation ref")
		}
		prepared = append(prepared, citation)
	}
	return prepared, nil
}

func validCitation(citation Citation) bool {
	switch citation.CitationKind {
	case CitationKindKnowledge:
		return citation.KnowledgeBaseID != "" && citation.RootKind != "" &&
			citation.RepoRole != "" && citation.Revision != "" &&
			citation.Path != "" && citation.LineStart >= 0 && citation.LineEnd >= citation.LineStart
	case CitationKindWeb:
		return citation.URL != "" && citation.LineStart >= 0 && citation.LineEnd >= citation.LineStart
	default:
		return false
	}
}

func insertCitations(
	ctx context.Context,
	tx coredb.Tx,
	workspaceID string,
	turnID string,
	attemptID string,
	citations []Citation,
	now time.Time,
) error {
	if len(citations) == 0 {
		return nil
	}
	values := make([]string, 0, len(citations))
	args := make([]any, 0, len(citations)*20)
	for _, citation := range citations {
		offset := len(args)
		values = append(values, fmt.Sprintf(
			"($%d,$%d,$%d,$%d,$%d,$%d,NULLIF($%d,''),$%d,$%d,$%d,$%d,$%d,$%d,$%d,$%d,NULLIF($%d,''),NULLIF($%d,''),NULLIF($%d,''),$%d,$%d)",
			offset+1, offset+2, offset+3, offset+4, offset+5, offset+6,
			offset+7, offset+8, offset+9, offset+10, offset+11, offset+12,
			offset+13, offset+14, offset+15, offset+16, offset+17, offset+18,
			offset+19, offset+20,
		))
		args = append(args,
			citation.ID, workspaceID, turnID, attemptID,
			citation.CitationKind, citation.KnowledgeBaseID, citation.SourceID,
			citation.RootKind, citation.RepoRole, citation.Revision,
			citation.ResolvedRef, citation.ResolvedRefKind, citation.Path,
			citation.LineStart, citation.LineEnd, citation.URL, citation.Title,
			citation.Snippet, citation.FetchedAt, now,
		)
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_citations (
  id, workspace_id, turn_id, attempt_id, citation_kind, knowledge_base_id,
  source_id, root_kind, repo_role, revision, resolved_ref, resolved_ref_kind,
  path, line_start, line_end, url, title, snippet, fetched_at, created_at
) VALUES `+strings.Join(values, ","), args...); err != nil {
		return fmt.Errorf("persist Knowledge Q&A citations: %w", err)
	}
	return nil
}