package knowledgeqa

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

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

type rowScanner interface{ Scan(...any) error }

func scanSession(row rowScanner) (Session, error) {
	var session Session
	var kbJSON, sourceJSON []byte
	err := row.Scan(&session.ID, &session.WorkspaceID, &session.OrganizationID,
		&session.OwnerUserID, &session.OwnerMemberID, &session.RootMessageID,
		&session.ThreadChannelID, &kbJSON, &sourceJSON, &session.RequestedModel,
		&session.RequestedReasonLevel, &session.WebSearchEnabled, &session.ActiveTurnID,
		&session.Status, &session.LastActivityAt, &session.CreatedAt,
		&session.UpdatedAt)
	if err != nil {
		if coredb.IsNotFound(err) {
			return Session{}, ErrNotFound
		}
		return Session{}, fmt.Errorf("load Knowledge Q&A session: %w", err)
	}
	if err := decodeStringList(kbJSON, &session.RequestedKnowledgeBaseIDs); err != nil {
		return Session{}, fmt.Errorf("decode requested Knowledge Q&A KB ids: %w", err)
	}
	if err := decodeStringList(sourceJSON, &session.RequestedSourceIDs); err != nil {
		return Session{}, fmt.Errorf("decode requested Knowledge Q&A source ids: %w", err)
	}
	return session, nil
}

func turnSelectSQL() string {
	return `
SELECT turn.id, turn.workspace_id, turn.organization_id, turn.session_id,
       turn.requester_user_id, turn.user_message_id,
       turn.assistant_message_id, turn.question, turn.idempotency_key,
       turn.requested_knowledge_base_ids_json,
       turn.requested_source_ids_json, turn.authorized_roots_json,
       turn.requested_model, turn.requested_reason_level,
       turn.web_search_enabled,
       turn.effective_model, turn.effective_reason_level,
       turn.organization_policy_version,
       turn.runtime_version, turn.prompt_version,
       turn.current_attempt_id, turn.status, turn.error_code,
       turn.error_detail, turn.queued_at, turn.started_at,
       turn.completed_at, turn.updated_at
`
}

func loadTurn(ctx context.Context, queryer coredb.QueryRower, workspaceID, turnID string) (Turn, error) {
	return scanTurn(queryer.QueryRow(ctx, turnSelectSQL()+`
FROM knowledge_qa_turns AS turn
WHERE turn.workspace_id=$1 AND turn.id=$2
`, workspaceID, turnID))
}

func scanTurn(row rowScanner) (Turn, error) {
	var turn Turn
	var kbJSON, sourceJSON, rootsJSON []byte
	err := row.Scan(&turn.ID, &turn.WorkspaceID, &turn.OrganizationID,
		&turn.SessionID, &turn.RequesterUserID, &turn.UserMessageID,
		&turn.AssistantMessageID, &turn.Question, &turn.IdempotencyKey,
		&kbJSON, &sourceJSON, &rootsJSON, &turn.RequestedModel,
		&turn.RequestedReasonLevel, &turn.WebSearchEnabled, &turn.EffectiveModel,
		&turn.EffectiveReasonLevel, &turn.OrganizationPolicyVersion, &turn.RuntimeVersion,
		&turn.PromptVersion, &turn.CurrentAttemptID, &turn.Status,
		&turn.ErrorCode, &turn.ErrorDetail, &turn.QueuedAt, &turn.StartedAt,
		&turn.CompletedAt, &turn.UpdatedAt)
	if err != nil {
		if coredb.IsNotFound(err) {
			return Turn{}, ErrNotFound
		}
		return Turn{}, fmt.Errorf("load Knowledge Q&A turn: %w", err)
	}
	if err := decodeStringList(kbJSON, &turn.RequestedKnowledgeBaseIDs); err != nil {
		return Turn{}, fmt.Errorf("decode requested Knowledge Q&A KB ids: %w", err)
	}
	if err := decodeStringList(sourceJSON, &turn.RequestedSourceIDs); err != nil {
		return Turn{}, fmt.Errorf("decode requested Knowledge Q&A source ids: %w", err)
	}
	if len(rootsJSON) > 0 {
		roots, err := unmarshalAuthorizedRoots(rootsJSON)
		if err != nil {
			return Turn{}, fmt.Errorf("decode authorized Knowledge Q&A roots: %w", err)
		}
		turn.AuthorizedRoots = roots
	}
	return turn, nil
}

func loadTurnByIdempotencyKey(ctx context.Context, queryer coredb.QueryRower, sessionID, key string) (Turn, bool, error) {
	turn, err := scanTurn(queryer.QueryRow(ctx, turnSelectSQL()+`
FROM knowledge_qa_turns AS turn
WHERE turn.session_id=$1 AND turn.idempotency_key=$2
`, sessionID, key))
	if errors.Is(err, ErrNotFound) {
		return Turn{}, false, nil
	}
	return turn, err == nil, err
}

func lockAdmissionPrincipals(ctx context.Context, tx coredb.Tx, organizationID, userID string) error {
	var found string
	if err := tx.QueryRow(ctx, `SELECT id FROM organizations WHERE id=$1 FOR UPDATE`, organizationID).Scan(&found); err != nil {
		return fmt.Errorf("lock Knowledge Q&A organization admission: %w", err)
	}
	if err := tx.QueryRow(ctx, `SELECT id FROM users WHERE id=$1 FOR UPDATE`, userID).Scan(&found); err != nil {
		return fmt.Errorf("lock Knowledge Q&A user admission: %w", err)
	}
	return nil
}

func enforceSubmissionAdmission(ctx context.Context, tx coredb.Tx, organizationID, userID string, now time.Time) error {
	var userQueued, userRunning, orgRunning, userRecent, orgRecent int
	windowStart := now.Add(-time.Minute)
	if err := tx.QueryRow(ctx, `
SELECT
  COALESCE(SUM(CASE WHEN requester_user_id=$2 AND status='queued' THEN 1 ELSE 0 END),0),
  COALESCE(SUM(CASE WHEN requester_user_id=$2 AND status IN ('preparing','generating') THEN 1 ELSE 0 END),0),
  COALESCE(SUM(CASE WHEN status IN ('preparing','generating') THEN 1 ELSE 0 END),0)
FROM knowledge_qa_turns WHERE organization_id=$1
`, organizationID, userID).Scan(&userQueued, &userRunning,
		&orgRunning); err != nil {
		return fmt.Errorf("load Knowledge Q&A admission counters: %w", err)
	}
	if err := tx.QueryRow(ctx, `
SELECT
  COALESCE(SUM(CASE WHEN turn.requester_user_id=$2 THEN 1 ELSE 0 END),0),
  COUNT(*)
FROM knowledge_qa_turn_attempts AS attempt
JOIN knowledge_qa_turns AS turn ON turn.id=attempt.turn_id
WHERE attempt.organization_id=$1 AND attempt.queued_at >= $3
`, organizationID, userID, windowStart).Scan(&userRecent, &orgRecent); err != nil {
		return fmt.Errorf("load Knowledge Q&A submission rate: %w", err)
	}
	if userRecent >= MaxUserSubmissionsPerMinute || orgRecent >= MaxOrgSubmissionsPerMinute {
		return ErrRateLimited
	}
	if userQueued >= MaxUserQueuedTurns {
		return ErrQueueFull
	}
	// Running capacity is enforced again by the worker when it atomically
	// claims a queued attempt. Capture impossible/corrupt over-capacity here.
	if userRunning > MaxUserRunningTurns || orgRunning > MaxOrganizationRunningTurns {
		return ErrQueueFull
	}
	return nil
}

func loadUserQuestion(ctx context.Context, tx coredb.Tx, session Session, userMessageID string) (string, error) {
	var question string
	if err := tx.QueryRow(ctx, `
SELECT content FROM scope_messages
WHERE workspace_id=$1 AND id=$2 AND author_type='human'
  AND author_member_id=$3 AND (scope_id=$4 OR id=$5)
`, session.WorkspaceID, userMessageID, session.OwnerMemberID,
		session.ThreadChannelID, session.RootMessageID).Scan(&question); err != nil {
		if coredb.IsNotFound(err) {
			return "", ErrForbidden
		}
		return "", fmt.Errorf("load Knowledge Q&A user message: %w", err)
	}
	question = strings.TrimSpace(question)
	if question == "" {
		return "", ErrForbidden
	}
	return question, nil
}