-- MPC Distributed Signature System Database Schema
-- Version: 001
-- Description: Initial schema creation

-- Enable UUID extension
CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- ============================================
-- Session Coordinator Schema
-- ============================================

-- MPC Sessions table
CREATE TABLE mpc_sessions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_type VARCHAR(20) NOT NULL,      -- 'keygen' or 'sign'
    threshold_n INTEGER NOT NULL,
    threshold_t INTEGER NOT NULL,
    status VARCHAR(20) NOT NULL,
    message_hash BYTEA,                     -- For Sign sessions
    public_key BYTEA,                       -- Group public key after Keygen completion
    created_by VARCHAR(255) NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP NOT NULL DEFAULT NOW(),
    expires_at TIMESTAMP NOT NULL,
    completed_at TIMESTAMP,
    CONSTRAINT chk_threshold CHECK (threshold_t <= threshold_n AND threshold_t > 0),
    CONSTRAINT chk_session_type CHECK (session_type IN ('keygen', 'sign')),
    CONSTRAINT chk_status CHECK (status IN ('created', 'in_progress', 'completed', 'failed', 'expired'))
);

-- Indexes for mpc_sessions
CREATE INDEX idx_mpc_sessions_status ON mpc_sessions(status);
CREATE INDEX idx_mpc_sessions_created_at ON mpc_sessions(created_at);
CREATE INDEX idx_mpc_sessions_expires_at ON mpc_sessions(expires_at);
CREATE INDEX idx_mpc_sessions_created_by ON mpc_sessions(created_by);

-- Session Participants table
CREATE TABLE participants (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID NOT NULL REFERENCES mpc_sessions(id) ON DELETE CASCADE,
    party_id VARCHAR(255) NOT NULL,
    party_index INTEGER NOT NULL,
    status VARCHAR(20) NOT NULL,
    device_type VARCHAR(50),
    device_id VARCHAR(255),
    platform VARCHAR(50),
    app_version VARCHAR(50),
    public_key BYTEA,                       -- Party identity public key (for authentication)
    joined_at TIMESTAMP NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMP,
    CONSTRAINT chk_participant_status CHECK (status IN ('invited', 'joined', 'ready', 'completed', 'failed')),
    UNIQUE(session_id, party_id),
    UNIQUE(session_id, party_index)
);

-- Indexes for participants
CREATE INDEX idx_participants_session_id ON participants(session_id);
CREATE INDEX idx_participants_party_id ON participants(party_id);
CREATE INDEX idx_participants_status ON participants(status);

-- ============================================
-- Message Router Schema
-- ============================================

-- MPC Messages table (for offline message caching)
CREATE TABLE mpc_messages (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    session_id UUID NOT NULL REFERENCES mpc_sessions(id) ON DELETE CASCADE,
    from_party VARCHAR(255) NOT NULL,
    to_parties TEXT[],                      -- NULL means broadcast
    round_number INTEGER NOT NULL,
    message_type VARCHAR(50) NOT NULL,
    payload BYTEA NOT NULL,                 -- Encrypted MPC message
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    delivered_at TIMESTAMP,
    CONSTRAINT chk_round_number CHECK (round_number >= 0)
);

-- Indexes for mpc_messages
CREATE INDEX idx_mpc_messages_session_id ON mpc_messages(session_id);
CREATE INDEX idx_mpc_messages_to_parties ON mpc_messages USING GIN(to_parties);
CREATE INDEX idx_mpc_messages_delivered_at ON mpc_messages(delivered_at) WHERE delivered_at IS NULL;
CREATE INDEX idx_mpc_messages_created_at ON mpc_messages(created_at);
CREATE INDEX idx_mpc_messages_round ON mpc_messages(session_id, round_number);

-- ============================================
-- Server Party Service Schema
-- ============================================

-- Party Key Shares table (Server Party's own Share)
CREATE TABLE party_key_shares (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    party_id VARCHAR(255) NOT NULL,
    party_index INTEGER NOT NULL,
    session_id UUID NOT NULL,               -- Keygen session ID
    threshold_n INTEGER NOT NULL,
    threshold_t INTEGER NOT NULL,
    share_data BYTEA NOT NULL,              -- Encrypted tss-lib LocalPartySaveData
    public_key BYTEA NOT NULL,              -- Group public key
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    last_used_at TIMESTAMP,
    CONSTRAINT chk_key_share_threshold CHECK (threshold_t <= threshold_n)
);

-- Indexes for party_key_shares
CREATE INDEX idx_party_key_shares_party_id ON party_key_shares(party_id);
CREATE INDEX idx_party_key_shares_session_id ON party_key_shares(session_id);
CREATE INDEX idx_party_key_shares_public_key ON party_key_shares(public_key);
CREATE UNIQUE INDEX idx_party_key_shares_unique ON party_key_shares(party_id, session_id);

-- ============================================
-- Account Service Schema
-- ============================================

-- Accounts table
-- Only username is required, all other fields are optional
CREATE TABLE accounts (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    username VARCHAR(255) UNIQUE NOT NULL,
    email VARCHAR(255),
    phone VARCHAR(50),
    public_key BYTEA,                       -- MPC group public key (optional, set after keygen)
    keygen_session_id UUID,                 -- Related Keygen session (optional)
    threshold_n INTEGER,
    threshold_t INTEGER,
    status VARCHAR(20) DEFAULT 'active',
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP NOT NULL DEFAULT NOW(),
    last_login_at TIMESTAMP,
    CONSTRAINT chk_account_status CHECK (status IN ('active', 'suspended', 'locked', 'recovering'))
);

-- Indexes for accounts
CREATE INDEX idx_accounts_username ON accounts(username);
CREATE INDEX idx_accounts_email ON accounts(email) WHERE email IS NOT NULL;
CREATE UNIQUE INDEX idx_accounts_email_unique ON accounts(email) WHERE email IS NOT NULL;
CREATE INDEX idx_accounts_public_key ON accounts(public_key);
CREATE INDEX idx_accounts_status ON accounts(status);

-- Account Share Mapping table (records share locations, not share content)
CREATE TABLE account_shares (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
    share_type VARCHAR(20) NOT NULL,        -- 'user_device', 'server', 'recovery'
    party_id VARCHAR(255) NOT NULL,
    party_index INTEGER NOT NULL,
    device_type VARCHAR(50),
    device_id VARCHAR(255),
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    last_used_at TIMESTAMP,
    is_active BOOLEAN DEFAULT TRUE,
    CONSTRAINT chk_share_type CHECK (share_type IN ('user_device', 'server', 'recovery'))
);

-- Indexes for account_shares
CREATE INDEX idx_account_shares_account_id ON account_shares(account_id);
CREATE INDEX idx_account_shares_party_id ON account_shares(party_id);
CREATE INDEX idx_account_shares_active ON account_shares(account_id, is_active) WHERE is_active = TRUE;

-- Account Recovery Sessions table
CREATE TABLE account_recovery_sessions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    account_id UUID NOT NULL REFERENCES accounts(id),
    recovery_type VARCHAR(20) NOT NULL,     -- 'device_lost', 'share_rotation'
    old_share_type VARCHAR(20),
    new_keygen_session_id UUID,
    status VARCHAR(20) NOT NULL,
    requested_at TIMESTAMP NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMP,
    CONSTRAINT chk_recovery_status CHECK (status IN ('requested', 'in_progress', 'completed', 'failed'))
);

-- Indexes for account_recovery_sessions
CREATE INDEX idx_account_recovery_account_id ON account_recovery_sessions(account_id);
CREATE INDEX idx_account_recovery_status ON account_recovery_sessions(status);

-- ============================================
-- Audit Service Schema
-- ============================================

-- Audit Workflows table
CREATE TABLE audit_workflows (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    workflow_name VARCHAR(255) NOT NULL,
    workflow_type VARCHAR(50) NOT NULL,
    data_hash BYTEA NOT NULL,
    threshold_n INTEGER NOT NULL,
    threshold_t INTEGER NOT NULL,
    sign_session_id UUID,                   -- Related signing session
    signature BYTEA,
    status VARCHAR(20) NOT NULL,
    created_by VARCHAR(255) NOT NULL,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP NOT NULL DEFAULT NOW(),
    expires_at TIMESTAMP,
    completed_at TIMESTAMP,
    metadata JSONB,
    CONSTRAINT chk_audit_workflow_status CHECK (status IN ('pending', 'in_progress', 'approved', 'rejected', 'expired'))
);

-- Indexes for audit_workflows
CREATE INDEX idx_audit_workflows_status ON audit_workflows(status);
CREATE INDEX idx_audit_workflows_created_at ON audit_workflows(created_at);
CREATE INDEX idx_audit_workflows_workflow_type ON audit_workflows(workflow_type);

-- Audit Approvers table
CREATE TABLE audit_approvers (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    workflow_id UUID NOT NULL REFERENCES audit_workflows(id) ON DELETE CASCADE,
    approver_id VARCHAR(255) NOT NULL,
    party_id VARCHAR(255) NOT NULL,
    party_index INTEGER NOT NULL,
    status VARCHAR(20) NOT NULL,
    approved_at TIMESTAMP,
    comments TEXT,
    CONSTRAINT chk_approver_status CHECK (status IN ('pending', 'approved', 'rejected')),
    UNIQUE(workflow_id, approver_id)
);

-- Indexes for audit_approvers
CREATE INDEX idx_audit_approvers_workflow_id ON audit_approvers(workflow_id);
CREATE INDEX idx_audit_approvers_approver_id ON audit_approvers(approver_id);
CREATE INDEX idx_audit_approvers_status ON audit_approvers(status);

-- ============================================
-- Shared Audit Logs Schema
-- ============================================

-- Audit Logs table (shared across all services)
CREATE TABLE audit_logs (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    service_name VARCHAR(100) NOT NULL,
    action_type VARCHAR(100) NOT NULL,
    user_id VARCHAR(255),
    resource_type VARCHAR(100),
    resource_id VARCHAR(255),
    session_id UUID,
    ip_address INET,
    user_agent TEXT,
    request_data JSONB,
    response_data JSONB,
    status VARCHAR(20) NOT NULL,
    error_message TEXT,
    created_at TIMESTAMP NOT NULL DEFAULT NOW(),
    CONSTRAINT chk_audit_status CHECK (status IN ('success', 'failure', 'pending'))
);

-- Indexes for audit_logs
CREATE INDEX idx_audit_logs_created_at ON audit_logs(created_at);
CREATE INDEX idx_audit_logs_user_id ON audit_logs(user_id);
CREATE INDEX idx_audit_logs_session_id ON audit_logs(session_id);
CREATE INDEX idx_audit_logs_action_type ON audit_logs(action_type);
CREATE INDEX idx_audit_logs_service_name ON audit_logs(service_name);

-- Partitioning for audit_logs (if needed for large scale)
-- CREATE TABLE audit_logs_y2024m01 PARTITION OF audit_logs
--     FOR VALUES FROM ('2024-01-01') TO ('2024-02-01');

-- ============================================
-- Helper Functions
-- ============================================

-- Function to update updated_at timestamp
CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = NOW();
    RETURN NEW;
END;
$$ language 'plpgsql';

-- Triggers for auto-updating updated_at
CREATE TRIGGER update_mpc_sessions_updated_at
    BEFORE UPDATE ON mpc_sessions
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_accounts_updated_at
    BEFORE UPDATE ON accounts
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER update_audit_workflows_updated_at
    BEFORE UPDATE ON audit_workflows
    FOR EACH ROW EXECUTE FUNCTION update_updated_at_column();

-- Function to cleanup expired sessions
CREATE OR REPLACE FUNCTION cleanup_expired_sessions()
RETURNS INTEGER AS $$
DECLARE
    deleted_count INTEGER;
BEGIN
    UPDATE mpc_sessions
    SET status = 'expired', updated_at = NOW()
    WHERE expires_at < NOW()
    AND status IN ('created', 'in_progress');

    GET DIAGNOSTICS deleted_count = ROW_COUNT;
    RETURN deleted_count;
END;
$$ language 'plpgsql';

-- Function to cleanup old messages
CREATE OR REPLACE FUNCTION cleanup_old_messages(retention_hours INTEGER DEFAULT 24)
RETURNS INTEGER AS $$
DECLARE
    deleted_count INTEGER;
BEGIN
    DELETE FROM mpc_messages
    WHERE created_at < NOW() - (retention_hours || ' hours')::INTERVAL;

    GET DIAGNOSTICS deleted_count = ROW_COUNT;
    RETURN deleted_count;
END;
$$ language 'plpgsql';

-- Comments
COMMENT ON TABLE mpc_sessions IS 'MPC session management - Coordinator does not participate in MPC computation';
COMMENT ON TABLE participants IS 'Session participants - tracks join status of each party';
COMMENT ON TABLE mpc_messages IS 'MPC protocol messages - encrypted, router does not decrypt';
COMMENT ON TABLE party_key_shares IS 'Server party key shares - encrypted storage of tss-lib data';
COMMENT ON TABLE accounts IS 'User accounts with MPC-based authentication';
COMMENT ON TABLE audit_logs IS 'Comprehensive audit trail for all operations';