-- Migration 004: User Management and Authentication
-- This migration creates comprehensive user management for HFT trading systems

-- Enable required extensions
CREATE EXTENSION IF NOT EXISTS "pgcrypto";

-- Users table - comprehensive user management
CREATE TABLE IF NOT EXISTS users (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    username VARCHAR(64) NOT NULL UNIQUE,
    email VARCHAR(255) NOT NULL UNIQUE,
    password_hash TEXT NOT NULL,
    salt TEXT NOT NULL,
    first_name VARCHAR(100),
    last_name VARCHAR(100),
    phone VARCHAR(20),
    time_zone VARCHAR(50) DEFAULT 'UTC',
    language VARCHAR(10) DEFAULT 'en',
    status VARCHAR(20) NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'inactive', 'suspended', 'locked')),
    role VARCHAR(50) NOT NULL DEFAULT 'trader' CHECK (role IN ('admin', 'trader', 'risk_manager', 'analyst', 'readonly')),
    last_login TIMESTAMP WITH TIME ZONE,
    failed_login_attempts INTEGER NOT NULL DEFAULT 0,
    locked_until TIMESTAMP WITH TIME ZONE,
    password_changed_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    password_expires_at TIMESTAMP WITH TIME ZONE,
    two_factor_enabled BOOLEAN NOT NULL DEFAULT false,
    two_factor_secret TEXT,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    created_by UUID REFERENCES users(id),
    updated_by UUID REFERENCES users(id),
    metadata JSONB
);

-- Sessions table - track user sessions for security
CREATE TABLE IF NOT EXISTS user_sessions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    session_token TEXT NOT NULL UNIQUE,
    refresh_token TEXT,
    ip_address INET NOT NULL,
    user_agent TEXT,
    location JSONB,
    is_active BOOLEAN NOT NULL DEFAULT true,
    expires_at TIMESTAMP WITH TIME ZONE NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    last_used_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB
);

-- API Keys table - for programmatic access
CREATE TABLE IF NOT EXISTS api_keys (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    key_name VARCHAR(100) NOT NULL,
    key_hash TEXT NOT NULL UNIQUE,
    key_prefix VARCHAR(20) NOT NULL,
    permissions JSONB NOT NULL DEFAULT '[]'::jsonb,
    rate_limit INTEGER DEFAULT 1000, -- requests per minute
    is_active BOOLEAN NOT NULL DEFAULT true,
    last_used_at TIMESTAMP WITH TIME ZONE,
    usage_count BIGINT NOT NULL DEFAULT 0,
    expires_at TIMESTAMP WITH TIME ZONE,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB
);

-- Accounts table - trading accounts linked to users
CREATE TABLE IF NOT EXISTS accounts (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    account_number VARCHAR(64) NOT NULL UNIQUE,
    account_name VARCHAR(100) NOT NULL,
    user_id UUID NOT NULL REFERENCES users(id),
    account_type VARCHAR(20) NOT NULL DEFAULT 'individual' CHECK (account_type IN ('individual', 'corporate', 'institutional', 'demo')),
    base_currency VARCHAR(10) NOT NULL DEFAULT 'USD',
    initial_balance DECIMAL(20, 8) NOT NULL DEFAULT 0,
    current_balance DECIMAL(20, 8) NOT NULL DEFAULT 0,
    available_balance DECIMAL(20, 8) NOT NULL DEFAULT 0,
    margin_balance DECIMAL(20, 8) NOT NULL DEFAULT 0,
    equity DECIMAL(20, 8) NOT NULL DEFAULT 0,
    free_margin DECIMAL(20, 8) NOT NULL DEFAULT 0,
    margin_level DECIMAL(10, 4) NOT NULL DEFAULT 0, -- Margin level percentage
    leverage DECIMAL(10, 2) NOT NULL DEFAULT 1.00,
    max_leverage DECIMAL(10, 2) NOT NULL DEFAULT 100.00,
    status VARCHAR(20) NOT NULL DEFAULT 'active' CHECK (status IN ('active', 'inactive', 'suspended', 'closed')),
    risk_profile VARCHAR(20) NOT NULL DEFAULT 'medium' CHECK (risk_profile IN ('conservative', 'medium', 'aggressive', 'high_frequency')),
    broker VARCHAR(100),
    broker_account_id VARCHAR(100),
    is_demo BOOLEAN NOT NULL DEFAULT false,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB
);

-- Account permissions - granular access control
CREATE TABLE IF NOT EXISTS account_permissions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
    permission_type VARCHAR(50) NOT NULL, -- 'read', 'trade', 'admin', 'risk_override'
    granted_by UUID NOT NULL REFERENCES users(id),
    granted_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    expires_at TIMESTAMP WITH TIME ZONE,
    is_active BOOLEAN NOT NULL DEFAULT true,
    metadata JSONB,
    
    UNIQUE(user_id, account_id, permission_type)
);

-- Brokers table - external broker connections
CREATE TABLE IF NOT EXISTS brokers (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    name VARCHAR(100) NOT NULL UNIQUE,
    broker_type VARCHAR(50) NOT NULL, -- 'mt4', 'mt5', 'ctrader', 'fix', 'rest'
    api_endpoint TEXT,
    fix_settings JSONB,
    connection_settings JSONB NOT NULL DEFAULT '{}'::jsonb,
    credentials_encrypted TEXT,
    is_active BOOLEAN NOT NULL DEFAULT true,
    is_demo BOOLEAN NOT NULL DEFAULT false,
    supported_symbols TEXT[], -- Array of supported symbols
    commission_settings JSONB,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB
);

-- Broker connections - track live connections
CREATE TABLE IF NOT EXISTS broker_connections (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    broker_id UUID NOT NULL REFERENCES brokers(id) ON DELETE CASCADE,
    account_id UUID NOT NULL REFERENCES accounts(id) ON DELETE CASCADE,
    connection_id VARCHAR(100) NOT NULL,
    status VARCHAR(20) NOT NULL DEFAULT 'disconnected' CHECK (status IN ('connected', 'connecting', 'disconnected', 'error')),
    last_heartbeat TIMESTAMP WITH TIME ZONE,
    latency_ms INTEGER,
    error_message TEXT,
    connection_attempts INTEGER NOT NULL DEFAULT 0,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB,
    
    UNIQUE(broker_id, account_id)
);

-- Compliance profiles - regulatory requirements
CREATE TABLE IF NOT EXISTS compliance_profiles (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    profile_name VARCHAR(100) NOT NULL UNIQUE,
    jurisdiction VARCHAR(10) NOT NULL, -- 'US', 'EU', 'UK', 'APAC'
    regulations JSONB NOT NULL DEFAULT '{}'::jsonb,
    requirements JSONB NOT NULL DEFAULT '{}'::jsonb,
    reporting_requirements JSONB NOT NULL DEFAULT '{}'::jsonb,
    retention_periods JSONB NOT NULL DEFAULT '{}'::jsonb,
    is_active BOOLEAN NOT NULL DEFAULT true,
    created_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    metadata JSONB
);

-- User compliance assignments
CREATE TABLE IF NOT EXISTS user_compliance (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE,
    compliance_profile_id UUID NOT NULL REFERENCES compliance_profiles(id),
    assigned_at TIMESTAMP WITH TIME ZONE NOT NULL DEFAULT NOW(),
    assigned_by UUID NOT NULL REFERENCES users(id),
    is_active BOOLEAN NOT NULL DEFAULT true,
    metadata JSONB,
    
    UNIQUE(user_id, compliance_profile_id)
);

-- Create optimized indexes for HFT performance

-- Users table indexes
CREATE INDEX IF NOT EXISTS idx_users_username ON users(username);
CREATE INDEX IF NOT EXISTS idx_users_email ON users(email);
CREATE INDEX IF NOT EXISTS idx_users_status ON users(status);
CREATE INDEX IF NOT EXISTS idx_users_role ON users(role);
CREATE INDEX IF NOT EXISTS idx_users_last_login ON users(last_login);

-- Sessions table indexes
CREATE INDEX IF NOT EXISTS idx_user_sessions_user_id ON user_sessions(user_id);
CREATE INDEX IF NOT EXISTS idx_user_sessions_token ON user_sessions(session_token);
CREATE INDEX IF NOT EXISTS idx_user_sessions_active ON user_sessions(is_active, expires_at);
CREATE INDEX IF NOT EXISTS idx_user_sessions_ip_address ON user_sessions(ip_address);

-- API Keys table indexes
CREATE INDEX IF NOT EXISTS idx_api_keys_user_id ON api_keys(user_id);
CREATE INDEX IF NOT EXISTS idx_api_keys_prefix ON api_keys(key_prefix);
CREATE INDEX IF NOT EXISTS idx_api_keys_active ON api_keys(is_active, expires_at);

-- Accounts table indexes
CREATE INDEX IF NOT EXISTS idx_accounts_user_id ON accounts(user_id);
CREATE INDEX IF NOT EXISTS idx_accounts_number ON accounts(account_number);
CREATE INDEX IF NOT EXISTS idx_accounts_status ON accounts(status);
CREATE INDEX IF NOT EXISTS idx_accounts_type ON accounts(account_type);
CREATE INDEX IF NOT EXISTS idx_accounts_broker ON accounts(broker);

-- Account permissions indexes
CREATE INDEX IF NOT EXISTS idx_account_permissions_user_account ON account_permissions(user_id, account_id);
CREATE INDEX IF NOT EXISTS idx_account_permissions_type ON account_permissions(permission_type);
CREATE INDEX IF NOT EXISTS idx_account_permissions_active ON account_permissions(is_active, expires_at);

-- Brokers table indexes
CREATE INDEX IF NOT EXISTS idx_brokers_name ON brokers(name);
CREATE INDEX IF NOT EXISTS idx_brokers_type ON brokers(broker_type);
CREATE INDEX IF NOT EXISTS idx_brokers_active ON brokers(is_active);

-- Broker connections indexes
CREATE INDEX IF NOT EXISTS idx_broker_connections_broker_account ON broker_connections(broker_id, account_id);
CREATE INDEX IF NOT EXISTS idx_broker_connections_status ON broker_connections(status);
CREATE INDEX IF NOT EXISTS idx_broker_connections_heartbeat ON broker_connections(last_heartbeat);

-- Create triggers for automatic updates
CREATE TRIGGER trigger_users_updated_at
    BEFORE UPDATE ON users
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();

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

CREATE TRIGGER trigger_brokers_updated_at
    BEFORE UPDATE ON brokers
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER trigger_broker_connections_updated_at
    BEFORE UPDATE ON broker_connections
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();

-- Create functions for user management

-- Function to create new user with encrypted password
CREATE OR REPLACE FUNCTION create_user(
    p_username VARCHAR(64),
    p_email VARCHAR(255),
    p_password TEXT,
    p_first_name VARCHAR(100) DEFAULT NULL,
    p_last_name VARCHAR(100) DEFAULT NULL,
    p_role VARCHAR(50) DEFAULT 'trader',
    p_created_by UUID DEFAULT NULL
)
RETURNS UUID AS $$
DECLARE
    v_user_id UUID;
    v_salt TEXT;
    v_password_hash TEXT;
BEGIN
    -- Generate salt and hash password
    v_salt := encode(gen_random_bytes(32), 'hex');
    v_password_hash := crypt(p_password || v_salt, gen_salt('bf', 12));
    
    -- Insert user
    INSERT INTO users (
        username, email, password_hash, salt, first_name, last_name, 
        role, created_by, password_changed_at
    ) VALUES (
        p_username, p_email, v_password_hash, v_salt, p_first_name, p_last_name,
        p_role, p_created_by, NOW()
    ) RETURNING id INTO v_user_id;
    
    -- Create audit log entry
    INSERT INTO audit_logs (
        event_type, entity_type, entity_id, user_id, action, 
        new_values, timestamp, source
    ) VALUES (
        'user_created', 'user', v_user_id, p_created_by, 'create',
        jsonb_build_object('username', p_username, 'email', p_email, 'role', p_role),
        NOW(), 'user_management'
    );
    
    RETURN v_user_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Function to authenticate user
CREATE OR REPLACE FUNCTION authenticate_user(
    p_username VARCHAR(64),
    p_password TEXT,
    p_ip_address INET DEFAULT NULL,
    p_user_agent TEXT DEFAULT NULL
)
RETURNS TABLE(
    user_id UUID,
    session_token TEXT,
    expires_at TIMESTAMP WITH TIME ZONE,
    role VARCHAR(50),
    status VARCHAR(20)
) AS $$
DECLARE
    v_user_record RECORD;
    v_session_token TEXT;
    v_expires_at TIMESTAMP WITH TIME ZONE;
BEGIN
    -- Get user record
    SELECT u.id, u.username, u.password_hash, u.salt, u.status, u.role, u.failed_login_attempts, u.locked_until
    INTO v_user_record
    FROM users u
    WHERE u.username = p_username OR u.email = p_username;
    
    -- Check if user exists
    IF v_user_record.id IS NULL THEN
        RAISE EXCEPTION 'Invalid credentials';
    END IF;
    
    -- Check if account is locked
    IF v_user_record.locked_until IS NOT NULL AND v_user_record.locked_until > NOW() THEN
        RAISE EXCEPTION 'Account is locked until %', v_user_record.locked_until;
    END IF;
    
    -- Check if account is active
    IF v_user_record.status != 'active' THEN
        RAISE EXCEPTION 'Account is not active';
    END IF;
    
    -- Verify password
    IF NOT (v_user_record.password_hash = crypt(p_password || v_user_record.salt, v_user_record.password_hash)) THEN
        -- Increment failed login attempts
        UPDATE users 
        SET failed_login_attempts = failed_login_attempts + 1,
            locked_until = CASE 
                WHEN failed_login_attempts >= 4 THEN NOW() + INTERVAL '30 minutes'
                ELSE NULL 
            END
        WHERE id = v_user_record.id;
        
        RAISE EXCEPTION 'Invalid credentials';
    END IF;
    
    -- Reset failed login attempts and update last login
    UPDATE users 
    SET failed_login_attempts = 0, 
        locked_until = NULL,
        last_login = NOW()
    WHERE id = v_user_record.id;
    
    -- Generate session token
    v_session_token := encode(gen_random_bytes(64), 'hex');
    v_expires_at := NOW() + INTERVAL '8 hours';
    
    -- Create session
    INSERT INTO user_sessions (
        user_id, session_token, ip_address, user_agent, expires_at, last_used_at
    ) VALUES (
        v_user_record.id, v_session_token, p_ip_address, p_user_agent, v_expires_at, NOW()
    );
    
    -- Log successful login
    INSERT INTO audit_logs (
        event_type, entity_type, entity_id, user_id, action,
        new_values, timestamp, source, ip_address, user_agent
    ) VALUES (
        'user_login', 'user', v_user_record.id, v_user_record.id, 'login',
        jsonb_build_object('ip_address', p_ip_address::TEXT),
        NOW(), 'authentication', p_ip_address, p_user_agent
    );
    
    -- Return session info
    RETURN QUERY SELECT 
        v_user_record.id, 
        v_session_token, 
        v_expires_at,
        v_user_record.role,
        v_user_record.status;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Function to validate session
CREATE OR REPLACE FUNCTION validate_session(p_session_token TEXT)
RETURNS TABLE(
    user_id UUID,
    username VARCHAR(64),
    role VARCHAR(50),
    expires_at TIMESTAMP WITH TIME ZONE,
    is_valid BOOLEAN
) AS $$
BEGIN
    -- Update last used time and return session info
    UPDATE user_sessions 
    SET last_used_at = NOW()
    WHERE session_token = p_session_token 
      AND is_active = true 
      AND expires_at > NOW();
    
    RETURN QUERY
    SELECT 
        u.id,
        u.username,
        u.role,
        s.expires_at,
        (s.id IS NOT NULL AND s.expires_at > NOW()) as is_valid
    FROM user_sessions s
    JOIN users u ON s.user_id = u.id
    WHERE s.session_token = p_session_token 
      AND s.is_active = true
      AND u.status = 'active';
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Function to create trading account
CREATE OR REPLACE FUNCTION create_trading_account(
    p_user_id UUID,
    p_account_name VARCHAR(100),
    p_account_type VARCHAR(20) DEFAULT 'individual',
    p_initial_balance DECIMAL(20, 8) DEFAULT 0,
    p_leverage DECIMAL(10, 2) DEFAULT 1.00,
    p_is_demo BOOLEAN DEFAULT false
)
RETURNS UUID AS $$
DECLARE
    v_account_id UUID;
    v_account_number VARCHAR(64);
BEGIN
    -- Generate account number
    v_account_number := 'AC' || to_char(NOW(), 'YYYYMMDD') || '-' || 
                       encode(gen_random_bytes(4), 'hex');
    
    -- Insert account
    INSERT INTO accounts (
        user_id, account_number, account_name, account_type,
        initial_balance, current_balance, available_balance,
        leverage, is_demo
    ) VALUES (
        p_user_id, v_account_number, p_account_name, p_account_type,
        p_initial_balance, p_initial_balance, p_initial_balance,
        p_leverage, p_is_demo
    ) RETURNING id INTO v_account_id;
    
    -- Grant full permissions to account owner
    INSERT INTO account_permissions (
        user_id, account_id, permission_type, granted_by
    ) VALUES 
        (p_user_id, v_account_id, 'read', p_user_id),
        (p_user_id, v_account_id, 'trade', p_user_id),
        (p_user_id, v_account_id, 'admin', p_user_id);
    
    -- Create audit log
    INSERT INTO audit_logs (
        event_type, entity_type, entity_id, user_id, action,
        new_values, timestamp, source
    ) VALUES (
        'account_created', 'account', v_account_id, p_user_id, 'create',
        jsonb_build_object(
            'account_number', v_account_number,
            'account_type', p_account_type,
            'initial_balance', p_initial_balance,
            'is_demo', p_is_demo
        ),
        NOW(), 'account_management'
    );
    
    RETURN v_account_id;
END;
$$ LANGUAGE plpgsql SECURITY DEFINER;

-- Add constraints for data integrity
ALTER TABLE users ADD CONSTRAINT check_password_expiry 
    CHECK (password_expires_at IS NULL OR password_expires_at > password_changed_at);

ALTER TABLE accounts ADD CONSTRAINT check_balances 
    CHECK (current_balance >= 0 AND available_balance >= 0);

ALTER TABLE accounts ADD CONSTRAINT check_leverage 
    CHECK (leverage > 0 AND leverage <= max_leverage);

-- Create views for common queries

-- Active users view
CREATE VIEW active_users AS
SELECT 
    id, username, email, first_name, last_name, role, 
    last_login, created_at, two_factor_enabled
FROM users 
WHERE status = 'active';

-- Account summary view
CREATE VIEW account_summary AS
SELECT 
    a.id, a.account_number, a.account_name, a.account_type,
    u.username, u.first_name, u.last_name,
    a.current_balance, a.available_balance, a.equity,
    a.leverage, a.status, a.is_demo,
    COUNT(p.id) as position_count,
    COUNT(o.id) as open_orders
FROM accounts a
JOIN users u ON a.user_id = u.id
LEFT JOIN positions p ON a.id::text = p.account_id AND p.quantity != 0
LEFT JOIN orders o ON a.id::text = o.account_id AND o.status IN ('pending', 'partial')
GROUP BY a.id, u.username, u.first_name, u.last_name;

-- Add comments for documentation
COMMENT ON TABLE users IS 'Comprehensive user management for HFT trading systems';
COMMENT ON TABLE accounts IS 'Trading accounts with real-time balance tracking';
COMMENT ON TABLE api_keys IS 'API keys for programmatic trading access';
COMMENT ON TABLE user_sessions IS 'Active user sessions for security tracking';
COMMENT ON TABLE brokers IS 'External broker connection configurations';
COMMENT ON TABLE compliance_profiles IS 'Regulatory compliance requirements';

COMMENT ON FUNCTION create_user IS 'Create new user with encrypted password and audit trail';
COMMENT ON FUNCTION authenticate_user IS 'Authenticate user and create session with security logging';
COMMENT ON FUNCTION validate_session IS 'Validate active session and update last used timestamp';
COMMENT ON FUNCTION create_trading_account IS 'Create new trading account with permissions';