package knowledgeqa

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

	"github.com/tidbcloud/AgentTeam/internal/core/identifiers"
	"github.com/tidbcloud/AgentTeam/internal/coredb"
	"github.com/tidbcloud/AgentTeam/internal/knowledgeqa/protocol"
	"github.com/tidbcloud/AgentTeam/internal/store/model"
)

func (store *Store) LoadPolicyMetrics(ctx context.Context, organizationID string) (protocol.PolicyMetrics, error) {
	if store == nil || store.db == nil {
		return protocol.PolicyMetrics{}, fmt.Errorf("knowledge Q&A policy store is not configured")
	}
	organizationID = strings.TrimSpace(organizationID)
	if organizationID == "" {
		return protocol.PolicyMetrics{}, fmt.Errorf("organization id is required")
	}
	const windowDays = 30
	metrics := protocol.PolicyMetrics{WindowDays: windowDays}
	if err := store.db.QueryRow(ctx, `
SELECT COUNT(*),
       COALESCE(SUM(CASE WHEN status='completed' THEN 1 ELSE 0 END),0),
       COALESCE(SUM(CASE WHEN status='no_answer' THEN 1 ELSE 0 END),0),
       COALESCE(SUM(CASE WHEN status='failed' THEN 1 ELSE 0 END),0),
       COALESCE(SUM(CASE WHEN status='canceled' THEN 1 ELSE 0 END),0),
       COALESCE(SUM(CASE WHEN status IN ('preparing','generating') THEN 1 ELSE 0 END),0),
       COALESCE(SUM(CASE WHEN status='queued' THEN 1 ELSE 0 END),0),
       COALESCE(AVG(CASE
         WHEN started_at IS NOT NULL AND completed_at IS NOT NULL
         THEN TIMESTAMPDIFF(MICROSECOND, started_at, completed_at) / 1000
       END),0)
FROM knowledge_qa_turns
WHERE organization_id=$1
  AND queued_at >= DATE_SUB(CURRENT_TIMESTAMP, INTERVAL 30 DAY)
`, organizationID).Scan(
		&metrics.TotalTurns, &metrics.CompletedTurns, &metrics.NoAnswerTurns,
		&metrics.FailedTurns, &metrics.CanceledTurns, &metrics.RunningTurns,
		&metrics.QueuedTurns, &metrics.AverageTerminalLatencyMS,
	); err != nil {
		return protocol.PolicyMetrics{}, fmt.Errorf("load Knowledge Q&A operational metrics: %w", err)
	}
	if err := store.loadPolicyUsageMetrics(ctx, organizationID, &metrics); err != nil {
		return protocol.PolicyMetrics{}, err
	}
	if err := store.loadPolicyFailureMetrics(ctx, organizationID, &metrics); err != nil {
		return protocol.PolicyMetrics{}, err
	}
	return metrics, nil
}

func (store *Store) loadPolicyUsageMetrics(ctx context.Context, organizationID string, metrics *protocol.PolicyMetrics) error {
	if err := store.db.QueryRow(ctx, `
SELECT COUNT(*),
       COALESCE(SUM(CAST(COALESCE(
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.input_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.prompt_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.inputTokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.promptTokens')), '0'
       ) AS SIGNED)),0),
       COALESCE(SUM(CAST(COALESCE(
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.cache_read_input_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.cached_input_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.prompt_tokens_details.cached_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.input_tokens_details.cached_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.cachedInputTokens')), '0'
       ) AS SIGNED)),0),
       COALESCE(SUM(CAST(COALESCE(
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.output_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.completion_tokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.outputTokens')),
         JSON_UNQUOTE(JSON_EXTRACT(usage_json,'$.completionTokens')), '0'
       ) AS SIGNED)),0)
FROM knowledge_qa_turn_attempts
WHERE organization_id=$1 AND settlement_ref <> ''
  AND queued_at >= DATE_SUB(CURRENT_TIMESTAMP, INTERVAL 30 DAY)
`, organizationID).Scan(
		&metrics.SettledUsageTurns, &metrics.InputTokens,
		&metrics.CachedInputTokens, &metrics.OutputTokens,
	); err != nil {
		return fmt.Errorf("load Knowledge Q&A usage metrics: %w", err)
	}
	return nil
}

func (store *Store) loadPolicyFailureMetrics(ctx context.Context, organizationID string, metrics *protocol.PolicyMetrics) error {
	rows, err := store.db.Query(ctx, `
SELECT COALESCE(NULLIF(error_code,''),'knowledge_qa_runtime_failed'), COUNT(*)
FROM knowledge_qa_turn_attempts
WHERE organization_id=$1 AND status='failed'
  AND queued_at >= DATE_SUB(CURRENT_TIMESTAMP, INTERVAL 30 DAY)
GROUP BY COALESCE(NULLIF(error_code,''),'knowledge_qa_runtime_failed')
ORDER BY COUNT(*) DESC, COALESCE(NULLIF(error_code,''),'knowledge_qa_runtime_failed') ASC
LIMIT 20
`, organizationID)
	if err != nil {
		return fmt.Errorf("load Knowledge Q&A failure metrics: %w", err)
	}
	defer rows.Close()
	metrics.FailureTypes = make([]protocol.PolicyFailureCount, 0)
	for rows.Next() {
		var item protocol.PolicyFailureCount
		if err := rows.Scan(&item.Code, &item.Count); err != nil {
			return fmt.Errorf("scan Knowledge Q&A failure metrics: %w", err)
		}
		metrics.FailureTypes = append(metrics.FailureTypes, item)
	}
	if err := rows.Err(); err != nil {
		return fmt.Errorf("iterate Knowledge Q&A failure metrics: %w", err)
	}
	return nil
}

func (store *Store) LoadOrganizationPolicy(ctx context.Context, organizationID string) (protocol.OrganizationPolicy, error) {
	if store == nil || store.db == nil {
		return protocol.OrganizationPolicy{}, fmt.Errorf("knowledge Q&A policy store is not configured")
	}
	organizationID = strings.TrimSpace(organizationID)
	if organizationID == "" {
		return protocol.OrganizationPolicy{}, fmt.Errorf("organization id is required")
	}
	return loadOrganizationPolicy(ctx, store.db, organizationID)
}

func (store *Store) UpdateOrganizationPolicy(ctx context.Context, actorUserID, organizationID string, update protocol.OrganizationPolicyUpdate, now time.Time) (protocol.OrganizationPolicy, error) {
	if store == nil || store.db == nil {
		return protocol.OrganizationPolicy{}, fmt.Errorf("knowledge Q&A policy store is not configured")
	}
	actorUserID = strings.TrimSpace(actorUserID)
	organizationID = strings.TrimSpace(organizationID)
	if actorUserID == "" || organizationID == "" {
		return protocol.OrganizationPolicy{}, fmt.Errorf("actor user id and organization id are required")
	}
	now = normalizeNow(now)
	var result protocol.OrganizationPolicy
	err := coredb.WithTx(ctx, store.db, func(tx coredb.Tx) error {
		var found string
		if err := tx.QueryRow(ctx, `SELECT id FROM organizations WHERE id=$1 FOR UPDATE`, organizationID).Scan(&found); err != nil {
			if coredb.IsNotFound(err) {
				return protocol.ErrPolicyUnavailable
			}
			return fmt.Errorf("lock Knowledge Q&A organization: %w", err)
		}
		current, err := loadOrganizationPolicy(ctx, tx, organizationID)
		if err != nil {
			return err
		}
		nextVersion := current.PolicyVersion + 1
		if _, err := tx.Exec(ctx, `
INSERT INTO organization_knowledge_qa_policy (
  organization_id, enabled, policy_version, updated_by_user_id, created_at, updated_at
) VALUES ($1,$2,$3,$4,$5,$5)
ON DUPLICATE KEY UPDATE enabled=VALUES(enabled), policy_version=VALUES(policy_version),
  updated_by_user_id=VALUES(updated_by_user_id), updated_at=VALUES(updated_at)
`, organizationID, update.Enabled, nextVersion, actorUserID, now); err != nil {
			return fmt.Errorf("update organization Knowledge Q&A policy: %w", err)
		}
		result, err = loadOrganizationPolicy(ctx, tx, organizationID)
		if err != nil {
			return err
		}
		return insertPolicyAudit(ctx, tx, organizationID, actorUserID, current, result, now)
	})
	return result, err
}

func (store *Store) ResolveEffectivePolicy(ctx context.Context, organizationID, model, reasonLevel, dataPolicyRequirementRef string) (protocol.EffectivePolicy, error) {
	if strings.TrimSpace(dataPolicyRequirementRef) != "" {
		return protocol.EffectivePolicy{}, protocol.ErrDataPolicyUnsupported
	}
	model, reasonLevel, err := protocol.NormalizeSelection(model, reasonLevel)
	if err != nil {
		return protocol.EffectivePolicy{}, err
	}
	policy, err := store.LoadOrganizationPolicy(ctx, organizationID)
	if err != nil {
		return protocol.EffectivePolicy{}, err
	}
	if !policy.Enabled {
		return protocol.EffectivePolicy{}, protocol.ErrOrganizationDisabled
	}
	return protocol.EffectivePolicy{
		OrganizationID:            policy.OrganizationID,
		Model:                     model,
		ReasonLevel:               reasonLevel,
		OrganizationPolicyVersion: policy.PolicyVersion,
		WorkloadKind:              "knowledge_qa",
	}, nil
}

func loadOrganizationPolicy(ctx context.Context, queryer coredb.Queryer, organizationID string) (protocol.OrganizationPolicy, error) {
	var policy protocol.OrganizationPolicy
	var createdAt, updatedAt *time.Time
	if err := queryer.QueryRow(ctx, `
SELECT organization.id, COALESCE(policy.enabled,FALSE),
       COALESCE(policy.policy_version,0), COALESCE(policy.updated_by_user_id,''),
       policy.created_at, policy.updated_at
FROM organizations AS organization
LEFT JOIN organization_knowledge_qa_policy AS policy ON policy.organization_id=organization.id
WHERE organization.id=$1
`, organizationID).Scan(&policy.OrganizationID, &policy.Enabled, &policy.PolicyVersion,
		&policy.UpdatedByUserID, &createdAt, &updatedAt); err != nil {
		if coredb.IsNotFound(err) {
			return protocol.OrganizationPolicy{}, protocol.ErrPolicyUnavailable
		}
		return protocol.OrganizationPolicy{}, fmt.Errorf("load organization Knowledge Q&A policy: %w", err)
	}
	if createdAt != nil {
		policy.CreatedAt = createdAt.UTC()
	}
	if updatedAt != nil {
		policy.UpdatedAt = updatedAt.UTC()
	}
	return policy, nil
}

func insertPolicyAudit(ctx context.Context, tx coredb.Tx, organizationID, actorUserID string, before, after any, now time.Time) error {
	id, err := identifiers.NewEntityID("kqaudit")
	if err != nil {
		return fmt.Errorf("generate Knowledge Q&A policy audit id: %w", err)
	}
	beforeJSON, err := model.MarshalJSONObject("Knowledge Q&A policy audit before state", before)
	if err != nil {
		return fmt.Errorf("encode Knowledge Q&A policy audit before state: %w", err)
	}
	afterJSON, err := model.MarshalJSONObject("Knowledge Q&A policy audit after state", after)
	if err != nil {
		return fmt.Errorf("encode Knowledge Q&A policy audit after state: %w", err)
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO knowledge_qa_policy_audit (
  id, organization_id, actor_user_id, action, before_json, after_json, created_at
) VALUES ($1,$2,$3,'organization_policy_updated',$4,$5,$6)
`, id, organizationID, actorUserID, beforeJSON, afterJSON, now); err != nil {
		return fmt.Errorf("insert Knowledge Q&A policy audit: %w", err)
	}
	return nil
}