package secrets

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"

	postgresplatform "github.com/example/remote-access-platform/backend/internal/platform/postgres"
)

var (
	ErrResourceSecretNotFound = errors.New("resource secret not found")
	ErrSecretAccessDenied     = errors.New("resource secret access denied")
	ErrSecretLeaseRequired    = errors.New("resource secret resolution requires lease proof")
)

type ResourceSecretStore struct {
	db        postgresplatform.DBTX
	encryptor *Encryptor
	now       func() time.Time
}

type ResourceSecretResolver interface {
	ResolveForSession(ctx context.Context, req ResolveResourceSecretRequest) (*ResolvedResourceSecret, error)
}

type ResourceSecretDescriptor struct {
	ID             string          `json:"id"`
	OrganizationID string          `json:"organization_id"`
	ResourceID     string          `json:"resource_id"`
	SecretRef      string          `json:"secret_ref"`
	Protocol       string          `json:"protocol"`
	Version        int             `json:"version"`
	KeyID          string          `json:"key_id"`
	Algorithm      string          `json:"algorithm"`
	Metadata       json.RawMessage `json:"metadata"`
	CreatedAt      time.Time       `json:"created_at"`
	RotatedAt      *time.Time      `json:"rotated_at,omitempty"`
}

type UpsertResourceSecretCommand struct {
	OrganizationID string
	ResourceID     string
	Protocol       string
	SecretRef      string
	Payload        json.RawMessage
	Metadata       json.RawMessage
	ActorUserID    string
}

type ResolveResourceSecretRequest struct {
	SecretRef      string
	OrganizationID string
	ResourceID     string
	SessionID      string
	WorkerID       string
	LeaseID        string
}

type ResolvedResourceSecret struct {
	Descriptor ResourceSecretDescriptor
	Payload    json.RawMessage
}

func NewResourceSecretStore(db postgresplatform.DBTX, encryptor *Encryptor) *ResourceSecretStore {
	return &ResourceSecretStore{db: db, encryptor: encryptor, now: time.Now}
}

func (s *ResourceSecretStore) WithDB(db postgresplatform.DBTX) *ResourceSecretStore {
	if s == nil {
		return nil
	}
	return &ResourceSecretStore{db: db, encryptor: s.encryptor, now: s.now}
}

func DefaultResourceSecretRef(organizationID, resourceID string) string {
	return "rap-secret://org/" + strings.TrimSpace(organizationID) + "/resources/" + strings.TrimSpace(resourceID) + "/primary"
}

func (s *ResourceSecretStore) Upsert(ctx context.Context, cmd UpsertResourceSecretCommand) (*ResourceSecretDescriptor, error) {
	if s == nil || s.encryptor == nil {
		return nil, ErrSecretEncryptionKeyMissing
	}
	payload, err := normalizeJSONObject(cmd.Payload)
	if err != nil {
		return nil, err
	}
	metadata, err := normalizeJSONObjectAllowEmpty(cmd.Metadata)
	if err != nil {
		return nil, err
	}
	secretRef := strings.TrimSpace(cmd.SecretRef)
	if secretRef == "" {
		secretRef = DefaultResourceSecretRef(cmd.OrganizationID, cmd.ResourceID)
	}
	protocol := strings.ToLower(strings.TrimSpace(cmd.Protocol))
	encrypted, err := s.encryptor.Encrypt(payload, ResourceSecretAAD(cmd.OrganizationID, cmd.ResourceID, secretRef, protocol))
	if err != nil {
		return nil, err
	}
	now := s.now().UTC()
	const query = `
INSERT INTO resource_secrets (
    organization_id, resource_id, secret_ref, protocol, version, key_id,
    algorithm, nonce, ciphertext, payload_sha256, metadata, created_by_user_id,
    created_at, rotated_at
) VALUES (
    $1::uuid, $2::uuid, $3, $4, 1, $5,
    $6, $7, $8, $9, $10::jsonb, NULLIF($11, '')::uuid,
    $12, NULL
)
ON CONFLICT (resource_id) DO UPDATE SET
    secret_ref = EXCLUDED.secret_ref,
    protocol = EXCLUDED.protocol,
    version = resource_secrets.version + 1,
    key_id = EXCLUDED.key_id,
    algorithm = EXCLUDED.algorithm,
    nonce = EXCLUDED.nonce,
    ciphertext = EXCLUDED.ciphertext,
    payload_sha256 = EXCLUDED.payload_sha256,
    metadata = EXCLUDED.metadata,
    created_by_user_id = EXCLUDED.created_by_user_id,
    rotated_at = EXCLUDED.created_at
RETURNING id::text, organization_id::text, resource_id::text, secret_ref,
          protocol, version, key_id, algorithm, metadata, created_at, rotated_at
`
	var descriptor ResourceSecretDescriptor
	if err := s.db.QueryRow(ctx, query,
		cmd.OrganizationID,
		cmd.ResourceID,
		secretRef,
		protocol,
		encrypted.KeyID,
		encrypted.Algorithm,
		encrypted.Nonce,
		encrypted.Ciphertext,
		encrypted.PayloadSHA256,
		metadata,
		cmd.ActorUserID,
		now,
	).Scan(
		&descriptor.ID,
		&descriptor.OrganizationID,
		&descriptor.ResourceID,
		&descriptor.SecretRef,
		&descriptor.Protocol,
		&descriptor.Version,
		&descriptor.KeyID,
		&descriptor.Algorithm,
		&descriptor.Metadata,
		&descriptor.CreatedAt,
		&descriptor.RotatedAt,
	); err != nil {
		return nil, fmt.Errorf("upsert resource secret: %w", err)
	}
	return &descriptor, nil
}

func (s *ResourceSecretStore) ResolveForSession(ctx context.Context, req ResolveResourceSecretRequest) (*ResolvedResourceSecret, error) {
	if s == nil || s.encryptor == nil {
		return nil, ErrSecretEncryptionKeyMissing
	}
	if strings.TrimSpace(req.LeaseID) == "" {
		return nil, ErrSecretLeaseRequired
	}
	const query = `
SELECT sec.id::text, sec.organization_id::text, sec.resource_id::text, sec.secret_ref,
       sec.protocol, sec.version, sec.key_id, sec.algorithm, sec.metadata,
       sec.created_at, sec.rotated_at, sec.nonce, sec.ciphertext,
       rs.organization_id::text, rs.resource_id::text, COALESCE(rs.worker_id, ''), rs.state
FROM resource_secrets sec
JOIN remote_sessions rs ON rs.resource_id = sec.resource_id
WHERE sec.secret_ref = $1 AND rs.id = $2::uuid
`
	var descriptor ResourceSecretDescriptor
	var nonce, ciphertext []byte
	var sessionOrganizationID, sessionResourceID, sessionWorkerID, sessionState string
	if err := s.db.QueryRow(ctx, query, req.SecretRef, req.SessionID).Scan(
		&descriptor.ID,
		&descriptor.OrganizationID,
		&descriptor.ResourceID,
		&descriptor.SecretRef,
		&descriptor.Protocol,
		&descriptor.Version,
		&descriptor.KeyID,
		&descriptor.Algorithm,
		&descriptor.Metadata,
		&descriptor.CreatedAt,
		&descriptor.RotatedAt,
		&nonce,
		&ciphertext,
		&sessionOrganizationID,
		&sessionResourceID,
		&sessionWorkerID,
		&sessionState,
	); err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, ErrResourceSecretNotFound
		}
		return nil, fmt.Errorf("resolve resource secret: %w", err)
	}
	if descriptor.OrganizationID != req.OrganizationID ||
		descriptor.ResourceID != req.ResourceID ||
		sessionOrganizationID != req.OrganizationID ||
		sessionResourceID != req.ResourceID ||
		sessionWorkerID != req.WorkerID ||
		!secretResolvableSessionState(sessionState) {
		return nil, ErrSecretAccessDenied
	}
	plaintext, err := s.encryptor.Decrypt(EncryptedPayload{
		Algorithm:  descriptor.Algorithm,
		KeyID:      descriptor.KeyID,
		Nonce:      nonce,
		Ciphertext: ciphertext,
	}, ResourceSecretAAD(descriptor.OrganizationID, descriptor.ResourceID, descriptor.SecretRef, descriptor.Protocol))
	if err != nil {
		return nil, err
	}
	return &ResolvedResourceSecret{
		Descriptor: descriptor,
		Payload:    json.RawMessage(plaintext),
	}, nil
}

func normalizeJSONObject(raw json.RawMessage) (json.RawMessage, error) {
	if len(raw) == 0 || !json.Valid(raw) {
		return nil, ErrSecretPayloadInvalid
	}
	var decoded map[string]any
	if err := json.Unmarshal(raw, &decoded); err != nil {
		return nil, ErrSecretPayloadInvalid
	}
	encoded, err := json.Marshal(decoded)
	if err != nil {
		return nil, err
	}
	return json.RawMessage(encoded), nil
}

func normalizeJSONObjectAllowEmpty(raw json.RawMessage) (json.RawMessage, error) {
	if len(raw) == 0 {
		return json.RawMessage(`{}`), nil
	}
	return normalizeJSONObject(raw)
}

func secretResolvableSessionState(state string) bool {
	switch state {
	case "starting", "active", "reconnecting":
		return true
	default:
		return false
	}
}