package store

import (
	"context"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"time"

	"github.com/jackc/pgx/v5"
)

const (
	securityEventTransactionLockTimeout = 5 * time.Second
	securityEventTransactionIdleTimeout = 15 * time.Second
)

func (s *Store) beginSecurityEventTransaction(ctx context.Context) (pgx.Tx, error) {
	tx, err := s.pool.Begin(ctx)
	if err != nil {
		return nil, err
	}
	settings := []struct {
		name  string
		value string
	}{
		{name: "idle_in_transaction_session_timeout", value: securityEventTransactionIdleTimeout.String()},
		{name: "lock_timeout", value: securityEventTransactionLockTimeout.String()},
	}
	for _, setting := range settings {
		if _, err := tx.Exec(ctx, `SELECT set_config($1, $2, true)`, setting.name, setting.value); err != nil {
			rollbackSecurityEventTransaction(tx)
			return nil, err
		}
	}
	return tx, nil
}

func rollbackSecurityEventTransaction(tx pgx.Tx) {
	rollbackTransaction(tx)
}

type ApplySessionRevokedInput struct {
	Issuer           string
	Audience         string
	JTI              string
	TransactionID    string
	SubjectIssuer    string
	ApplicationID    string
	SubjectType      string
	TenantID         string
	Subject          string
	EventTimestamp   time.Time
	InitiatingEntity string
}

type ApplySecurityEventResult struct {
	Duplicate       bool
	WatermarkMoved  bool
	SessionsDeleted int64
	AuditID         string
}

type SecurityEventEvaluation struct {
	Revoked              bool
	RequireIntrospection bool
	Mode                 string
}

func (s *Store) ApplySessionRevoked(ctx context.Context, input ApplySessionRevokedInput) (ApplySecurityEventResult, error) {
	if input.SubjectType == "" {
		input.SubjectType = "principal"
	}
	tx, err := s.beginSecurityEventTransaction(ctx)
	if err != nil {
		return ApplySecurityEventResult{}, err
	}
	defer rollbackSecurityEventTransaction(tx)
	subjectHash := shortSecurityEventHash(input.Subject)
	tag, err := tx.Exec(ctx, `
INSERT INTO gateway_security_event_receipts (
  issuer, jti, audience, event_type, transaction_id, tenant_id, subject_hash, event_timestamp,
  application_id, subject_type
)
VALUES ($1, $2::uuid, $3, $4, NULLIF($5, ''), $6, $7, $8, NULLIF($9, ''), $10)
ON CONFLICT (issuer, jti) DO NOTHING`,
		input.Issuer, input.JTI, input.Audience,
		"https://schemas.openid.net/secevent/caep/event-type/session-revoked",
		input.TransactionID, input.TenantID, subjectHash, input.EventTimestamp,
		input.ApplicationID, input.SubjectType,
	)
	if err != nil {
		return ApplySecurityEventResult{}, err
	}
	if tag.RowsAffected() == 0 {
		if err := tx.Commit(ctx); err != nil {
			return ApplySecurityEventResult{}, err
		}
		return ApplySecurityEventResult{Duplicate: true}, nil
	}

	var revokedAt time.Time
	err = tx.QueryRow(ctx, `
INSERT INTO gateway_oidc_revocation_watermarks (
  issuer, application_id, tenant_id, subject_type, subject, revoked_at, source_jti
)
VALUES ($1, $2, $3, $4, $5, $6, $7::uuid)
ON CONFLICT (issuer, application_id, tenant_id, subject_type, subject) DO UPDATE
SET revoked_at = EXCLUDED.revoked_at, source_jti = EXCLUDED.source_jti, updated_at = now()
WHERE EXCLUDED.revoked_at > gateway_oidc_revocation_watermarks.revoked_at
RETURNING revoked_at`, input.SubjectIssuer, input.ApplicationID, input.TenantID, input.SubjectType,
		input.Subject, input.EventTimestamp, input.JTI).Scan(&revokedAt)
	watermarkMoved := err == nil
	if err != nil && !errors.Is(err, pgx.ErrNoRows) {
		return ApplySecurityEventResult{}, err
	}
	var sessionsDeleted int64
	if watermarkMoved {
		switch {
		case input.ApplicationID != "" && input.SubjectType == "principal":
			tag, err = tx.Exec(ctx, `
DELETE FROM gateway_oidc_sessions session
USING gateway_oidc_user_bindings user_binding,gateway_oidc_tenant_bindings tenant_binding
WHERE session.oidc_user_binding_id=user_binding.id
  AND user_binding.tenant_binding_id=tenant_binding.id
  AND tenant_binding.issuer=$1
  AND tenant_binding.application_id=$2
  AND tenant_binding.external_tenant_id=$3
  AND user_binding.subject=$4`, input.SubjectIssuer, input.ApplicationID, input.TenantID, input.Subject)
		case input.ApplicationID != "" && input.SubjectType == "tenant":
			tag, err = tx.Exec(ctx, `
DELETE FROM gateway_oidc_sessions session
USING gateway_oidc_user_bindings user_binding,gateway_oidc_tenant_bindings tenant_binding
WHERE session.oidc_user_binding_id=user_binding.id
  AND user_binding.tenant_binding_id=tenant_binding.id
  AND tenant_binding.issuer=$1
  AND tenant_binding.application_id=$2
  AND tenant_binding.external_tenant_id=$3`, input.SubjectIssuer, input.ApplicationID, input.TenantID)
			if err == nil {
				_, err = tx.Exec(ctx, `UPDATE gateway_oidc_tenant_bindings
					SET access_status='disabled',updated_at=now()
					WHERE issuer=$1 AND application_id=$2 AND external_tenant_id=$3`,
					input.SubjectIssuer, input.ApplicationID, input.TenantID)
			}
		default:
			tag, err = tx.Exec(ctx, `
DELETE FROM gateway_oidc_sessions session
USING gateway_users gateway_user
WHERE session.gateway_user_id = gateway_user.id
  AND gateway_user.source = 'oidc'
  AND gateway_user.external_user_id = $1
  AND gateway_user.tenant_id = $2`, input.Subject, input.TenantID)
		}
		if err != nil {
			return ApplySecurityEventResult{}, err
		}
		sessionsDeleted = tag.RowsAffected()
	}
	jtiHash := shortSecurityEventHash(input.JTI)
	audit, err := s.RecordAuditLogTx(ctx, tx, AuditLogInput{
		Category: "identity", Action: "identity.oidc_session.revoked", ActorSource: "ssf",
		TargetType: "oidc_subject", TargetID: subjectHash,
		AfterState: map[string]any{"watermarkMoved": watermarkMoved, "sessionsDeleted": sessionsDeleted},
		Metadata: map[string]any{
			"issuer": input.SubjectIssuer, "transmitterIssuer": input.Issuer, "tenantId": input.TenantID, "jtiHash": jtiHash,
			"applicationId": input.ApplicationID, "subjectType": input.SubjectType, "initiatingEntity": input.InitiatingEntity,
		},
	})
	if err != nil {
		return ApplySecurityEventResult{}, err
	}
	if err := tx.Commit(ctx); err != nil {
		return ApplySecurityEventResult{}, err
	}
	return ApplySecurityEventResult{WatermarkMoved: watermarkMoved, SessionsDeleted: sessionsDeleted, AuditID: audit.ID}, nil
}

func (s *Store) EnsureSecurityEventStreamState(ctx context.Context, issuer, audience, streamID string) error {
	tx, err := s.beginSecurityEventTransaction(ctx)
	if err != nil {
		return err
	}
	defer rollbackSecurityEventTransaction(tx)
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_security_event_stream_state(issuer,audience,stream_id,mode)
VALUES($1,$2,$3::uuid,'bootstrap')
ON CONFLICT(issuer,audience) DO NOTHING`, issuer, audience, streamID); err != nil {
		return err
	}
	var currentStreamID string
	if err := tx.QueryRow(ctx, `SELECT stream_id::text FROM gateway_security_event_stream_state
		WHERE issuer=$1 AND audience=$2 FOR UPDATE`, issuer, audience).Scan(&currentStreamID); err != nil {
		return err
	}
	if currentStreamID != streamID {
		if _, err := tx.Exec(ctx, `DELETE FROM gateway_security_event_verification_challenges
			WHERE issuer=$1 AND audience=$2`, issuer, audience); err != nil {
			return err
		}
		if _, err := tx.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET stream_id=$3::uuid,mode='bootstrap',stream_status='unknown',
    last_verification_at=NULL,bootstrap_until=NULL,consecutive_verifications=0,
    pending_state_hash=NULL,pending_state_created_at=NULL,fallback_since=NULL,
    last_error_category=NULL,created_at=now(),updated_at=now()
WHERE issuer=$1 AND audience=$2`, issuer, audience, streamID); err != nil {
			return err
		}
	}
	return tx.Commit(ctx)
}

func (s *Store) BeginSecurityEventVerification(ctx context.Context, issuer, audience string, stateHash []byte, now time.Time) error {
	if len(stateHash) != sha256.Size {
		return errors.New("verification state hash must be SHA-256")
	}
	tx, err := s.beginSecurityEventTransaction(ctx)
	if err != nil {
		return err
	}
	defer rollbackSecurityEventTransaction(tx)
	tag, err := tx.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET pending_state_hash=$3,pending_state_created_at=$4,updated_at=now()
WHERE issuer=$1 AND audience=$2`, issuer, audience, stateHash, now)
	if err != nil {
		return err
	}
	if tag.RowsAffected() == 0 {
		return errors.New("security event stream state is missing")
	}
	if _, err := tx.Exec(ctx, `
INSERT INTO gateway_security_event_verification_challenges(issuer,audience,state_hash,created_at,expires_at)
VALUES($1,$2,$3,$4::timestamptz,$4::timestamptz + interval '180 seconds')
ON CONFLICT(issuer,audience,state_hash) DO NOTHING`, issuer, audience, stateHash, now); err != nil {
		return err
	}
	if _, err := tx.Exec(ctx, `DELETE FROM gateway_security_event_verification_challenges WHERE expires_at < $1`, now); err != nil {
		return err
	}
	return tx.Commit(ctx)
}

func (s *Store) ConfirmSecurityEventVerification(ctx context.Context, issuer, audience, streamID, jti string, stateHash []byte, now time.Time) (bool, error) {
	tx, err := s.beginSecurityEventTransaction(ctx)
	if err != nil {
		return false, err
	}
	defer rollbackSecurityEventTransaction(tx)
	tag, err := tx.Exec(ctx, `
INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash)
VALUES($1,$2::uuid,$3,$4,$5)
ON CONFLICT(issuer,jti) DO NOTHING`, issuer, jti, audience,
		"https://schemas.openid.net/secevent/ssf/event-type/verification", shortSecurityEventHash(streamID))
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() == 0 {
		return false, tx.Commit(ctx)
	}
	tag, err = tx.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET last_verification_at=$5,
    bootstrap_until=CASE WHEN last_verification_at IS NULL THEN $5 + interval '360 seconds' ELSE bootstrap_until END,
    consecutive_verifications=CASE
      WHEN stream_status <> 'enabled' THEN 0
      WHEN mode='introspection_fallback' THEN consecutive_verifications+1
      ELSE 1
    END,
    mode=CASE
      WHEN stream_status <> 'enabled' THEN mode
      WHEN mode='introspection_fallback' AND consecutive_verifications+1 >= 2
        AND (bootstrap_until IS NULL OR $5 >= bootstrap_until) THEN 'push_healthy'
      WHEN mode='introspection_fallback' THEN 'introspection_fallback'
      WHEN mode='push_healthy' THEN 'push_healthy'
      ELSE 'bootstrap'
    END,
    fallback_since=CASE
      WHEN stream_status='enabled' AND mode='introspection_fallback' AND consecutive_verifications+1 >= 2
        AND (bootstrap_until IS NULL OR $5 >= bootstrap_until) THEN NULL
      ELSE fallback_since
    END,
    pending_state_hash=NULL,pending_state_created_at=NULL,
    last_error_category=CASE WHEN stream_status='enabled' THEN NULL ELSE 'stream_' || stream_status END,
    updated_at=now()
WHERE issuer=$1 AND audience=$2 AND stream_id=$3::uuid
	AND EXISTS(SELECT 1 FROM gateway_security_event_verification_challenges challenge
	  WHERE challenge.issuer=$1 AND challenge.audience=$2 AND challenge.state_hash=$4 AND challenge.expires_at >= $5)`,
		issuer, audience, streamID, stateHash, now)
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() == 0 {
		return false, errors.New("invalid verification state")
	}
	if _, err := tx.Exec(ctx, `DELETE FROM gateway_security_event_verification_challenges
		WHERE issuer=$1 AND audience=$2 AND state_hash=$3`, issuer, audience, stateHash); err != nil {
		return false, err
	}
	if _, err := tx.Exec(ctx, `
UPDATE gateway_security_event_connections connection
SET lifecycle_status='enabled',last_error_category=NULL,version=version+1,updated_at=now()
WHERE connection.transmitter_issuer=$1 AND connection.audience=$2
  AND connection.lifecycle_status='degraded'
  AND EXISTS(
    SELECT 1 FROM gateway_security_event_stream_state state
    WHERE state.issuer=$1 AND state.audience=$2
      AND state.mode='push_healthy' AND state.stream_status='enabled'
  )`, issuer, audience); err != nil {
		return false, err
	}
	return true, tx.Commit(ctx)
}

func (s *Store) RecordSecurityEventReceipt(ctx context.Context, issuer, audience, jti, eventType, subject string) (bool, error) {
	tag, err := s.pool.Exec(ctx, `
INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash)
VALUES($1,$2::uuid,$3,$4,$5)
ON CONFLICT(issuer,jti) DO NOTHING`, issuer, jti, audience, eventType, shortSecurityEventHash(subject))
	return tag.RowsAffected() == 1, err
}

// ApplySecurityEventStreamUpdated records a transmitter-originated stream
// status change and immediately moves OIDC traffic to introspection. Push is
// trusted again only after the normal consecutive-verification recovery path.
func (s *Store) ApplySecurityEventStreamUpdated(ctx context.Context, issuer, audience, jti, streamID, status, reason string, now time.Time) (bool, error) {
	tx, err := s.beginSecurityEventTransaction(ctx)
	if err != nil {
		return false, err
	}
	defer rollbackSecurityEventTransaction(tx)
	tag, err := tx.Exec(ctx, `
INSERT INTO gateway_security_event_receipts(issuer,jti,audience,event_type,subject_hash)
VALUES($1,$2::uuid,$3,$4,$5)
ON CONFLICT(issuer,jti) DO NOTHING`, issuer, jti, audience,
		"https://schemas.openid.net/secevent/ssf/event-type/stream-updated", shortSecurityEventHash(streamID))
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() == 0 {
		return false, tx.Commit(ctx)
	}
	tag, err = tx.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET stream_status=$5,mode='introspection_fallback',fallback_since=COALESCE(fallback_since,$4),
    consecutive_verifications=0,last_error_category=$6,updated_at=now()
WHERE issuer=$1 AND audience=$2 AND stream_id=$3::uuid`, issuer, audience, streamID, now, status, "stream_"+status)
	if err != nil {
		return false, err
	}
	if tag.RowsAffected() == 0 {
		return false, errors.New("security event stream state is missing")
	}
	if _, err := s.RecordAuditLogTx(ctx, tx, AuditLogInput{
		Category: "identity", Action: "identity.ssf_stream.updated", ActorSource: "ssf",
		TargetType: "ssf_stream", TargetID: shortSecurityEventHash(streamID),
		AfterState: map[string]any{"status": status},
		Metadata:   map[string]any{"issuer": issuer, "audience": audience, "reasonCategory": shortSecurityEventHash(reason)},
	}); err != nil {
		return false, err
	}
	return true, tx.Commit(ctx)
}

func (s *Store) RecordSecurityEventHeartbeatFailure(ctx context.Context, issuer, audience, category string) error {
	_, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state
		SET last_error_category=$3,updated_at=now() WHERE issuer=$1 AND audience=$2`, issuer, audience, category)
	return err
}

func (s *Store) AdvanceSecurityEventStreamState(ctx context.Context, issuer, audience string, now time.Time, staleAfter time.Duration) error {
	if staleAfter <= 0 {
		return errors.New("security event stale threshold must be positive")
	}
	_, err := s.pool.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET mode='introspection_fallback',fallback_since=COALESCE(fallback_since,$3),
    consecutive_verifications=0,last_error_category='verification_stale',updated_at=now()
WHERE issuer=$1 AND audience=$2 AND mode <> 'introspection_fallback'
  AND COALESCE(last_verification_at,created_at) < $3::timestamptz - make_interval(secs => $4::double precision)`,
		issuer, audience, now, staleAfter.Seconds())
	return err
}

func (s *Store) EvaluateOIDCSecurityEvent(ctx context.Context, transmitterIssuer, audience, subjectIssuer, applicationID, tenantID, subject string, issuedAt, now time.Time, staleAfter time.Duration) (SecurityEventEvaluation, error) {
	var mode, streamStatus string
	var lastVerification, bootstrapUntil *time.Time
	var createdAt time.Time
	var revokedAt *time.Time
	err := s.pool.QueryRow(ctx, `
SELECT state.mode,state.stream_status,state.last_verification_at,state.bootstrap_until,state.created_at,watermark.revoked_at
FROM gateway_security_event_stream_state state
LEFT JOIN LATERAL (
  SELECT max(revoked_at) AS revoked_at
  FROM gateway_oidc_revocation_watermarks
  WHERE issuer=$3 AND application_id=$4 AND tenant_id=$5
    AND (
      (subject_type='principal' AND subject=$6)
      OR (subject_type='tenant' AND subject=$5)
    )
) watermark ON true
WHERE state.issuer=$1 AND state.audience=$2`,
		transmitterIssuer, audience, subjectIssuer, applicationID, tenantID, subject).Scan(
		&mode, &streamStatus, &lastVerification, &bootstrapUntil, &createdAt, &revokedAt,
	)
	if err != nil {
		return SecurityEventEvaluation{}, err
	}
	evaluation := SecurityEventEvaluation{Mode: mode, Revoked: revokedAt != nil && !issuedAt.After(*revokedAt)}
	stale := (lastVerification != nil && now.Sub(*lastVerification) > staleAfter) ||
		(lastVerification == nil && (mode != "bootstrap" || now.Sub(createdAt) > staleAfter))
	if stale {
		evaluation.Mode, evaluation.RequireIntrospection = "introspection_fallback", true
		_, err := s.pool.Exec(ctx, `
UPDATE gateway_security_event_stream_state
SET mode='introspection_fallback',fallback_since=COALESCE(fallback_since,$3),consecutive_verifications=0,
    last_error_category='verification_stale',updated_at=now()
WHERE issuer=$1 AND audience=$2 AND mode <> 'introspection_fallback'`, transmitterIssuer, audience, now)
		return evaluation, err
	}
	if mode == "bootstrap" {
		if streamStatus != "enabled" || bootstrapUntil == nil || now.Before(*bootstrapUntil) {
			evaluation.RequireIntrospection = true
			return evaluation, nil
		}
		evaluation.Mode = "push_healthy"
		_, err := s.pool.Exec(ctx, `UPDATE gateway_security_event_stream_state SET mode='push_healthy',updated_at=now()
WHERE issuer=$1 AND audience=$2 AND mode='bootstrap'`, transmitterIssuer, audience)
		return evaluation, err
	}
	evaluation.RequireIntrospection = mode == "introspection_fallback"
	return evaluation, nil
}

func (s *Store) SecurityEventMetrics(ctx context.Context, issuer, audience string) (string, *time.Time, error) {
	var mode string
	var lastVerificationAt *time.Time
	err := s.pool.QueryRow(ctx, `
SELECT mode,last_verification_at
FROM gateway_security_event_stream_state
WHERE issuer=$1 AND audience=$2`, issuer, audience).Scan(&mode, &lastVerificationAt)
	return mode, lastVerificationAt, err
}

func shortSecurityEventHash(value string) string {
	digest := sha256.Sum256([]byte(value))
	return hex.EncodeToString(digest[:])[:16]
}
