package postgres

import (
	"context"
	"database/sql"
	"encoding/json"
	"time"

	"github.com/google/uuid"
	"github.com/lib/pq"
	"github.com/rwadurian/mpc-system/pkg/logger"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/entities"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/repositories"
	"github.com/rwadurian/mpc-system/services/session-coordinator/domain/value_objects"
	"go.uber.org/zap"
)

// SessionPostgresRepo implements SessionRepository for PostgreSQL
type SessionPostgresRepo struct {
	db *sql.DB
}

// NewSessionPostgresRepo creates a new PostgreSQL session repository
func NewSessionPostgresRepo(db *sql.DB) *SessionPostgresRepo {
	return &SessionPostgresRepo{db: db}
}

// Save persists or updates a session (upsert)
func (r *SessionPostgresRepo) Save(ctx context.Context, session *entities.MPCSession) error {
	tx, err := r.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// Upsert session (insert or update on conflict)
	_, err = tx.ExecContext(ctx, `
		INSERT INTO mpc_sessions (
			id, session_type, threshold_n, threshold_t, status,
			message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			wallet_name, invite_code
		) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12, $13, $14, $15, $16, $17)
		ON CONFLICT (id) DO UPDATE SET
			status = EXCLUDED.status,
			message_hash = EXCLUDED.message_hash,
			public_key = EXCLUDED.public_key,
			delegate_party_id = EXCLUDED.delegate_party_id,
			keygen_session_id = EXCLUDED.keygen_session_id,
			updated_at = EXCLUDED.updated_at,
			completed_at = EXCLUDED.completed_at,
			version = EXCLUDED.version,
			wallet_name = EXCLUDED.wallet_name,
			invite_code = EXCLUDED.invite_code
	`,
		session.ID.UUID(),
		string(session.SessionType),
		session.Threshold.N(),
		session.Threshold.T(),
		session.Status.String(),
		session.MessageHash,
		session.PublicKey,
		session.DelegatePartyID,
		nullUUID(session.KeygenSessionID),
		session.CreatedBy,
		session.CreatedAt,
		session.UpdatedAt,
		session.ExpiresAt,
		session.CompletedAt,
		session.Version,
		session.WalletName,
		session.InviteCode,
	)
	if err != nil {
		return err
	}

	// Delete existing participants before inserting new ones
	_, err = tx.ExecContext(ctx, `DELETE FROM participants WHERE session_id = $1`, session.ID.UUID())
	if err != nil {
		return err
	}

	// Insert participants
	for _, p := range session.Participants {
		deviceInfoJSON, err := json.Marshal(p.DeviceInfo)
		if err != nil {
			return err
		}

		_, err = tx.ExecContext(ctx, `
			INSERT INTO participants (
				id, session_id, party_id, party_index, status,
				device_type, device_id, platform, app_version, public_key, joined_at, completed_at
			) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
		`,
			uuid.New(),
			session.ID.UUID(),
			p.PartyID.String(),
			p.PartyIndex,
			p.Status.String(),
			p.DeviceInfo.DeviceType,
			p.DeviceInfo.DeviceID,
			p.DeviceInfo.Platform,
			p.DeviceInfo.AppVersion,
			p.PublicKey,
			p.JoinedAt,
			p.CompletedAt,
		)
		if err != nil {
			return err
		}
		_ = deviceInfoJSON // Unused but could be stored as JSON
	}

	return tx.Commit()
}

// FindByID retrieves a session by SessionID
func (r *SessionPostgresRepo) FindByID(ctx context.Context, id value_objects.SessionID) (*entities.MPCSession, error) {
	return r.FindByUUID(ctx, id.UUID())
}

// FindByUUID retrieves a session by UUID
func (r *SessionPostgresRepo) FindByUUID(ctx context.Context, id uuid.UUID) (*entities.MPCSession, error) {
	var session sessionRow
	err := r.db.QueryRowContext(ctx, `
		SELECT id, session_type, threshold_n, threshold_t, status,
			   message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			   COALESCE(wallet_name, ''), COALESCE(invite_code, '')
		FROM mpc_sessions WHERE id = $1
	`, id).Scan(
		&session.ID,
		&session.SessionType,
		&session.ThresholdN,
		&session.ThresholdT,
		&session.Status,
		&session.MessageHash,
		&session.PublicKey,
		&session.DelegatePartyID,
		&session.KeygenSessionID,
		&session.CreatedBy,
		&session.CreatedAt,
		&session.UpdatedAt,
		&session.ExpiresAt,
		&session.CompletedAt,
		&session.Version,
		&session.WalletName,
		&session.InviteCode,
	)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, entities.ErrSessionNotFound
		}
		return nil, err
	}

	// Load participants
	participants, err := r.loadParticipants(ctx, id)
	if err != nil {
		return nil, err
	}

	// Convert nullable KeygenSessionID
	var keygenSessionID uuid.UUID
	if session.KeygenSessionID != nil {
		keygenSessionID = *session.KeygenSessionID
	}

	return entities.ReconstructSessionWithOptions(
		session.ID,
		session.SessionType,
		session.ThresholdT,
		session.ThresholdN,
		session.Status,
		session.MessageHash,
		session.PublicKey,
		session.DelegatePartyID,
		keygenSessionID,
		session.CreatedBy,
		session.CreatedAt,
		session.UpdatedAt,
		session.ExpiresAt,
		session.CompletedAt,
		participants,
		session.Version,
		&entities.ReconstructSessionOptions{
			WalletName: session.WalletName,
			InviteCode: session.InviteCode,
		},
	)
}

// FindByStatus retrieves sessions by status
func (r *SessionPostgresRepo) FindByStatus(ctx context.Context, status value_objects.SessionStatus) ([]*entities.MPCSession, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, session_type, threshold_n, threshold_t, status,
			   message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			   COALESCE(wallet_name, ''), COALESCE(invite_code, '')
		FROM mpc_sessions WHERE status = $1
	`, status.String())
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return r.scanSessions(ctx, rows)
}

// FindExpired retrieves all expired but not yet marked sessions
// Note: expires_at is stored in UTC, so we compare with NOW() AT TIME ZONE 'UTC'
func (r *SessionPostgresRepo) FindExpired(ctx context.Context) ([]*entities.MPCSession, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, session_type, threshold_n, threshold_t, status,
			   message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			   COALESCE(wallet_name, ''), COALESCE(invite_code, '')
		FROM mpc_sessions
		WHERE expires_at < (NOW() AT TIME ZONE 'UTC') AND status IN ('created', 'in_progress')
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return r.scanSessions(ctx, rows)
}

// FindActive retrieves all active sessions (created or in_progress)
func (r *SessionPostgresRepo) FindActive(ctx context.Context) ([]*entities.MPCSession, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, session_type, threshold_n, threshold_t, status,
			   message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			   COALESCE(wallet_name, ''), COALESCE(invite_code, '')
		FROM mpc_sessions
		WHERE status IN ('created', 'in_progress')
		ORDER BY created_at ASC
	`)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return r.scanSessions(ctx, rows)
}

// FindByCreator retrieves sessions created by a user
func (r *SessionPostgresRepo) FindByCreator(ctx context.Context, creatorID string) ([]*entities.MPCSession, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT id, session_type, threshold_n, threshold_t, status,
			   message_hash, public_key, delegate_party_id, keygen_session_id, created_by, created_at, updated_at, expires_at, completed_at, version,
			   COALESCE(wallet_name, ''), COALESCE(invite_code, '')
		FROM mpc_sessions WHERE created_by = $1
		ORDER BY created_at DESC
	`, creatorID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return r.scanSessions(ctx, rows)
}

// FindActiveByParticipant retrieves active sessions for a participant
func (r *SessionPostgresRepo) FindActiveByParticipant(ctx context.Context, partyID value_objects.PartyID) ([]*entities.MPCSession, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT s.id, s.session_type, s.threshold_n, s.threshold_t, s.status,
			   s.message_hash, s.public_key, s.delegate_party_id, s.keygen_session_id, s.created_by, s.created_at, s.updated_at, s.expires_at, s.completed_at, s.version,
			   COALESCE(s.wallet_name, ''), COALESCE(s.invite_code, '')
		FROM mpc_sessions s
		JOIN participants p ON s.id = p.session_id
		WHERE p.party_id = $1 AND s.status IN ('created', 'in_progress')
		ORDER BY s.created_at DESC
	`, partyID.String())
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	return r.scanSessions(ctx, rows)
}

// Update updates an existing session using optimistic locking
func (r *SessionPostgresRepo) Update(ctx context.Context, session *entities.MPCSession) error {
	tx, err := r.db.BeginTx(ctx, nil)
	if err != nil {
		return err
	}
	defer tx.Rollback()

	// Update session with optimistic locking (check version)
	// Version is incremented automatically to prevent concurrent modifications
	result, err := tx.ExecContext(ctx, `
		UPDATE mpc_sessions SET
			status = $1,
			public_key = $2,
			updated_at = $3,
			completed_at = $4,
			version = version + 1
		WHERE id = $5 AND version = $6
	`,
		session.Status.String(),
		session.PublicKey,
		session.UpdatedAt,
		session.CompletedAt,
		session.ID.UUID(),
		session.Version,
	)
	if err != nil {
		return err
	}

	// Check if the update affected any rows
	rowsAffected, err := result.RowsAffected()
	if err != nil {
		return err
	}

	// If no rows were affected, version mismatch occurred (concurrent modification)
	if rowsAffected == 0 {
		logger.Warn("optimistic lock conflict detected",
			zap.String("session_id", session.ID.UUID().String()),
			zap.Int64("expected_version", session.Version))
		return entities.ErrOptimisticLockConflict
	}

	// Update session version in memory to reflect the database state
	session.Version++

	// DEBUG: Log all participant statuses before update
	var participantStatuses []string
	for _, p := range session.Participants {
		participantStatuses = append(participantStatuses, p.PartyID.String()+":"+p.Status.String())
	}
	logger.Debug("repo.Update: updating participants",
		zap.String("session_id", session.ID.UUID().String()),
		zap.Strings("participant_statuses", participantStatuses))

	// Update each participant individually
	for _, p := range session.Participants {
		// DEBUG: Log before updating each participant
		logger.Debug("repo.Update: updating individual participant",
			zap.String("session_id", session.ID.UUID().String()),
			zap.String("party_id", p.PartyID.String()),
			zap.String("new_status", p.Status.String()))

		// Try UPDATE first
		result, err := tx.ExecContext(ctx, `
			UPDATE participants SET
				status = $1,
				public_key = $2,
				completed_at = $3,
				device_type = $4,
				device_id = $5,
				platform = $6,
				app_version = $7
			WHERE session_id = $8 AND party_id = $9
		`,
			p.Status.String(),
			p.PublicKey,
			p.CompletedAt,
			p.DeviceInfo.DeviceType,
			p.DeviceInfo.DeviceID,
			p.DeviceInfo.Platform,
			p.DeviceInfo.AppVersion,
			session.ID.UUID(),
			p.PartyID.String(),
		)
		if err != nil {
			return err
		}

		// If no rows affected, participant doesn't exist yet, INSERT it
		rowsAffected, err := result.RowsAffected()
		if err != nil {
			return err
		}

		// DEBUG: Log rows affected
		logger.Debug("repo.Update: participant update result",
			zap.String("session_id", session.ID.UUID().String()),
			zap.String("party_id", p.PartyID.String()),
			zap.Int64("rows_affected", rowsAffected))

		if rowsAffected == 0 {
			// Participant doesn't exist, INSERT it
			logger.Debug("repo.Update: inserting new participant",
				zap.String("session_id", session.ID.UUID().String()),
				zap.String("party_id", p.PartyID.String()))

			_, err = tx.ExecContext(ctx, `
				INSERT INTO participants (
					id, session_id, party_id, party_index, status,
					device_type, device_id, platform, app_version, public_key, joined_at, completed_at
				) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
			`,
				uuid.New(),
				session.ID.UUID(),
				p.PartyID.String(),
				p.PartyIndex,
				p.Status.String(),
				p.DeviceInfo.DeviceType,
				p.DeviceInfo.DeviceID,
				p.DeviceInfo.Platform,
				p.DeviceInfo.AppVersion,
				p.PublicKey,
				p.JoinedAt,
				p.CompletedAt,
			)
			if err != nil {
				return err
			}
		}
	}

	logger.Debug("repo.Update: committing transaction",
		zap.String("session_id", session.ID.UUID().String()))

	return tx.Commit()
}

// Delete removes a session
func (r *SessionPostgresRepo) Delete(ctx context.Context, id value_objects.SessionID) error {
	_, err := r.db.ExecContext(ctx, `DELETE FROM mpc_sessions WHERE id = $1`, id.UUID())
	return err
}

// DeleteExpired removes all expired sessions
func (r *SessionPostgresRepo) DeleteExpired(ctx context.Context) (int64, error) {
	result, err := r.db.ExecContext(ctx, `
		DELETE FROM mpc_sessions
		WHERE status = 'expired' AND expires_at < NOW() - INTERVAL '24 hours'
	`)
	if err != nil {
		return 0, err
	}
	return result.RowsAffected()
}

// Count returns the total number of sessions
func (r *SessionPostgresRepo) Count(ctx context.Context) (int64, error) {
	var count int64
	err := r.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM mpc_sessions`).Scan(&count)
	return count, err
}

// CountByStatus returns the number of sessions by status
func (r *SessionPostgresRepo) CountByStatus(ctx context.Context, status value_objects.SessionStatus) (int64, error) {
	var count int64
	err := r.db.QueryRowContext(ctx, `SELECT COUNT(*) FROM mpc_sessions WHERE status = $1`, status.String()).Scan(&count)
	return count, err
}

// Helper methods

func (r *SessionPostgresRepo) loadParticipants(ctx context.Context, sessionID uuid.UUID) ([]*entities.Participant, error) {
	rows, err := r.db.QueryContext(ctx, `
		SELECT party_id, party_index, status, device_type, device_id, platform, app_version, public_key, joined_at, completed_at
		FROM participants WHERE session_id = $1
		ORDER BY party_index
	`, sessionID)
	if err != nil {
		return nil, err
	}
	defer rows.Close()

	var participants []*entities.Participant
	for rows.Next() {
		var p participantRow
		err := rows.Scan(
			&p.PartyID,
			&p.PartyIndex,
			&p.Status,
			&p.DeviceType,
			&p.DeviceID,
			&p.Platform,
			&p.AppVersion,
			&p.PublicKey,
			&p.JoinedAt,
			&p.CompletedAt,
		)
		if err != nil {
			return nil, err
		}

		partyID, _ := value_objects.NewPartyID(p.PartyID)
		participant := &entities.Participant{
			PartyID:    partyID,
			PartyIndex: p.PartyIndex,
			Status:     value_objects.ParticipantStatus(p.Status),
			DeviceInfo: entities.DeviceInfo{
				DeviceType: entities.DeviceType(p.DeviceType),
				DeviceID:   p.DeviceID,
				Platform:   p.Platform,
				AppVersion: p.AppVersion,
			},
			PublicKey:   p.PublicKey,
			JoinedAt:    p.JoinedAt,
			CompletedAt: p.CompletedAt,
		}
		participants = append(participants, participant)
	}

	return participants, rows.Err()
}

func (r *SessionPostgresRepo) scanSessions(ctx context.Context, rows *sql.Rows) ([]*entities.MPCSession, error) {
	var sessions []*entities.MPCSession
	for rows.Next() {
		var s sessionRow
		err := rows.Scan(
			&s.ID,
			&s.SessionType,
			&s.ThresholdN,
			&s.ThresholdT,
			&s.Status,
			&s.MessageHash,
			&s.PublicKey,
			&s.DelegatePartyID,
			&s.KeygenSessionID,
			&s.CreatedBy,
			&s.CreatedAt,
			&s.UpdatedAt,
			&s.ExpiresAt,
			&s.CompletedAt,
			&s.Version,
			&s.WalletName,
			&s.InviteCode,
		)
		if err != nil {
			return nil, err
		}

		participants, err := r.loadParticipants(ctx, s.ID)
		if err != nil {
			return nil, err
		}

		// Convert nullable KeygenSessionID
		var keygenSessionID uuid.UUID
		if s.KeygenSessionID != nil {
			keygenSessionID = *s.KeygenSessionID
		}

		session, err := entities.ReconstructSessionWithOptions(
			s.ID,
			s.SessionType,
			s.ThresholdT,
			s.ThresholdN,
			s.Status,
			s.MessageHash,
			s.PublicKey,
			s.DelegatePartyID,
			keygenSessionID,
			s.CreatedBy,
			s.CreatedAt,
			s.UpdatedAt,
			s.ExpiresAt,
			s.CompletedAt,
			participants,
			s.Version,
			&entities.ReconstructSessionOptions{
				WalletName: s.WalletName,
				InviteCode: s.InviteCode,
			},
		)
		if err != nil {
			return nil, err
		}
		sessions = append(sessions, session)
	}

	return sessions, rows.Err()
}

// Row types for scanning
type sessionRow struct {
	ID              uuid.UUID
	SessionType     string
	ThresholdN      int
	ThresholdT      int
	Status          string
	MessageHash     []byte
	PublicKey       []byte
	DelegatePartyID string
	KeygenSessionID *uuid.UUID
	CreatedBy       string
	CreatedAt       time.Time
	UpdatedAt       time.Time
	ExpiresAt       time.Time
	CompletedAt     *time.Time
	Version         int64
	WalletName      string // 钱包名称 (for co_managed_keygen)
	InviteCode      string // 邀请码 (for co_managed_keygen)
}

type participantRow struct {
	PartyID     string
	PartyIndex  int
	Status      string
	DeviceType  string
	DeviceID    string
	Platform    string
	AppVersion  string
	PublicKey   []byte
	JoinedAt    time.Time
	CompletedAt *time.Time
}

// Ensure interface compliance
var _ repositories.SessionRepository = (*SessionPostgresRepo)(nil)

// Use pq for array handling
var _ = pq.Array

// nullUUID converts a UUID to a nullable database value
func nullUUID(u uuid.UUID) interface{} {
	if u == uuid.Nil {
		return nil
	}
	return u
}
