CREATE EXTENSION IF NOT EXISTS pgcrypto;

CREATE TYPE user_role AS ENUM (
    'admin',
    'operator',
    'viewer'
);

CREATE TYPE agent_status AS ENUM (
    'online',
    'offline',
    'busy',
    'error'
);

CREATE TYPE command_status AS ENUM (
    'pending',
    'running',
    'completed',
    'failed',
    'cancelled'
);

CREATE TYPE transfer_direction AS ENUM (
    'upload',
    'download'
);

CREATE TYPE actor_type_enum AS ENUM (
    'user',
    'agent',
    'system'
);

-- USERS

CREATE TABLE users (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    username TEXT NOT NULL UNIQUE,
    password_hash TEXT NOT NULL,

    role user_role NOT NULL DEFAULT 'viewer',

    is_active BOOLEAN NOT NULL DEFAULT TRUE,

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_users_role
ON users(role);

CREATE INDEX idx_users_active
ON users(is_active);

-- AGENTS

CREATE TABLE agents (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    hostname TEXT NOT NULL UNIQUE,

    token_hash TEXT NOT NULL,

    cert_fingerprint TEXT,

    os_info JSONB,

    status agent_status NOT NULL DEFAULT 'offline',

    is_active BOOLEAN NOT NULL DEFAULT TRUE,

    last_seen TIMESTAMPTZ,

    registered_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_agents_status
ON agents(status);

CREATE INDEX idx_agents_last_seen
ON agents(last_seen);

CREATE INDEX idx_agents_active
ON agents(is_active);

CREATE INDEX idx_agents_os_info_gin
ON agents
USING GIN(os_info);

-- COMMAND TYPES

CREATE TABLE command_types (
    id SMALLSERIAL PRIMARY KEY,
    name TEXT NOT NULL UNIQUE
);

INSERT INTO command_types(name)
VALUES
('shell'),
('powershell'),
('upload'),
('download'),
('system'),
('custom');

-- COMMANDS

CREATE TABLE commands (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    agent_id UUID NOT NULL
        REFERENCES agents(id)
        ON DELETE CASCADE,

    issued_by UUID
        REFERENCES users(id)
        ON DELETE SET NULL,

    command_type_id SMALLINT
        REFERENCES command_types(id),

    command TEXT NOT NULL,

    status command_status NOT NULL DEFAULT 'pending',

    output TEXT,

    exit_code INTEGER,

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    completed_at TIMESTAMPTZ,

    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    CHECK (
        completed_at IS NULL
        OR completed_at >= created_at
    ),

    CHECK (
        exit_code IS NULL
        OR exit_code BETWEEN -255 AND 255
    )
);

CREATE INDEX idx_commands_agent
ON commands(agent_id);

CREATE INDEX idx_commands_issued_by
ON commands(issued_by);

CREATE INDEX idx_commands_status
ON commands(status);

CREATE INDEX idx_commands_created_at
ON commands(created_at DESC);

CREATE INDEX idx_commands_pending
ON commands(status)
WHERE status = 'pending';

CREATE INDEX idx_commands_running
ON commands(status)
WHERE status = 'running';

-- FILE TRANSFERS

CREATE TABLE file_transfers (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    agent_id UUID
        REFERENCES agents(id)
        ON DELETE CASCADE,

    initiated_by UUID
        REFERENCES users(id)
        ON DELETE SET NULL,

    command_id UUID
        REFERENCES commands(id)
        ON DELETE SET NULL,

    direction transfer_direction NOT NULL,

    filename TEXT NOT NULL,

    sha256 CHAR(64),

    size_bytes BIGINT,

    status command_status NOT NULL DEFAULT 'pending',

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    CHECK (
        size_bytes IS NULL
        OR size_bytes >= 0
    )
);

CREATE INDEX idx_file_transfers_agent
ON file_transfers(agent_id);

CREATE INDEX idx_file_transfers_status
ON file_transfers(status);

CREATE INDEX idx_file_transfers_sha256
ON file_transfers(sha256);

CREATE INDEX idx_file_transfers_created
ON file_transfers(created_at DESC);

-- AUDIT LOGS

CREATE TABLE audit_logs (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    actor_id UUID,

    actor_type actor_type_enum NOT NULL,

    action TEXT NOT NULL,

    target TEXT,

    details JSONB,

    ip_address INET,

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_audit_actor
ON audit_logs(actor_id);

CREATE INDEX idx_audit_action
ON audit_logs(action);

CREATE INDEX idx_audit_created
ON audit_logs(created_at DESC);

CREATE INDEX idx_audit_details_gin
ON audit_logs
USING GIN(details);

-- AGENT HEARTBEATS

CREATE TABLE agent_heartbeats (
    id BIGSERIAL PRIMARY KEY,

    agent_id UUID NOT NULL
        REFERENCES agents(id)
        ON DELETE CASCADE,

    received_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),

    metadata JSONB
);

CREATE INDEX idx_heartbeats_agent
ON agent_heartbeats(agent_id);

CREATE INDEX idx_heartbeats_received
ON agent_heartbeats(received_at DESC);

-- USER SESSIONS

CREATE TABLE user_sessions (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    user_id UUID NOT NULL
        REFERENCES users(id)
        ON DELETE CASCADE,

    refresh_token_hash TEXT NOT NULL,

    user_agent TEXT,

    ip_address INET,

    expires_at TIMESTAMPTZ NOT NULL,

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE INDEX idx_sessions_user
ON user_sessions(user_id);

CREATE INDEX idx_sessions_expires
ON user_sessions(expires_at);

-- AGENT GROUPS

CREATE TABLE agent_groups (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),

    name TEXT NOT NULL UNIQUE,

    description TEXT,

    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

CREATE TABLE agent_group_members (
    agent_id UUID NOT NULL
        REFERENCES agents(id)
        ON DELETE CASCADE,

    group_id UUID NOT NULL
        REFERENCES agent_groups(id)
        ON DELETE CASCADE,

    PRIMARY KEY (agent_id, group_id)
);

-- UPDATE updated_at FUNCTION

CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS
$$
BEGIN
    NEW.updated_at = NOW();
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- updated_at TRIGGERS

CREATE TRIGGER trg_users_updated_at
BEFORE UPDATE ON users
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER trg_agents_updated_at
BEFORE UPDATE ON agents
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER trg_commands_updated_at
BEFORE UPDATE ON commands
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();

CREATE TRIGGER trg_file_transfers_updated_at
BEFORE UPDATE ON file_transfers
FOR EACH ROW
EXECUTE FUNCTION update_updated_at_column();

-- COMMAND COMPLETION TRIGGER

CREATE OR REPLACE FUNCTION update_agent_after_command()
RETURNS TRIGGER AS
$$
BEGIN
    IF NEW.status = 'completed'
       AND OLD.status IS DISTINCT FROM 'completed'
    THEN
        UPDATE agents
        SET
            last_seen = NOW(),
            status = 'online'
        WHERE id = NEW.agent_id;
    END IF;

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trg_update_agent_after_command
AFTER UPDATE ON commands
FOR EACH ROW
EXECUTE FUNCTION update_agent_after_command();

-- AUDIT LOG COMMAND CREATION

CREATE OR REPLACE FUNCTION log_command_creation()
RETURNS TRIGGER AS
$$
BEGIN
    INSERT INTO audit_logs (
        actor_id,
        actor_type,
        action,
        target,
        details
    )
    VALUES (
        NEW.issued_by,
        'user',
        'command_created',
        NEW.id::TEXT,
        jsonb_build_object(
            'command_length', length(NEW.command),
            'status', NEW.status
        )
    );

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trg_log_command_creation
AFTER INSERT ON commands
FOR EACH ROW
EXECUTE FUNCTION log_command_creation();

-- AGENT STATUS CHANGE LOGGING

CREATE OR REPLACE FUNCTION log_agent_status_change()
RETURNS TRIGGER AS
$$
BEGIN
    IF OLD.status IS DISTINCT FROM NEW.status THEN

        INSERT INTO audit_logs (
            actor_id,
            actor_type,
            action,
            target,
            details
        )
        VALUES (
            NEW.id,
            'agent',
            'status_changed',
            NEW.hostname,
            jsonb_build_object(
                'old_status', OLD.status,
                'new_status', NEW.status
            )
        );

    END IF;

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trg_log_agent_status_change
AFTER UPDATE ON agents
FOR EACH ROW
EXECUTE FUNCTION log_agent_status_change();


-- IMMUTABLE AUDIT LOGS

CREATE OR REPLACE FUNCTION prevent_audit_log_modification()
RETURNS TRIGGER AS
$$
BEGIN
    RAISE EXCEPTION 'audit_logs table is immutable';
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER trg_prevent_audit_update
BEFORE UPDATE ON audit_logs
FOR EACH ROW
EXECUTE FUNCTION prevent_audit_log_modification();

CREATE TRIGGER trg_prevent_audit_delete
BEFORE DELETE ON audit_logs
FOR EACH ROW
EXECUTE FUNCTION prevent_audit_log_modification();

-- VIEWS

CREATE VIEW active_agents_view AS
SELECT
    a.id,
    a.hostname,
    a.status,
    a.last_seen,
    a.registered_at,

    COUNT(c.id) AS total_commands,

    COUNT(*) FILTER (
        WHERE c.status = 'completed'
    ) AS completed_commands,

    COUNT(*) FILTER (
        WHERE c.status = 'pending'
    ) AS pending_commands

FROM agents a
LEFT JOIN commands c
ON a.id = c.agent_id

WHERE a.is_active = TRUE

GROUP BY
    a.id,
    a.hostname,
    a.status,
    a.last_seen,
    a.registered_at;


-- RECENT COMMANDS VIEW

CREATE VIEW recent_commands_view AS
SELECT
    c.id,
    c.command,
    c.status,
    c.created_at,
    c.completed_at,

    a.hostname AS agent_hostname,

    u.username AS issued_by

FROM commands c

JOIN agents a
ON c.agent_id = a.id

LEFT JOIN users u
ON c.issued_by = u.id;

-- AUDIT LOGS VIEW

CREATE VIEW audit_logs_view AS
SELECT
    al.id,
    al.action,
    al.target,
    al.created_at,
    al.ip_address,

    u.username AS actor_username,

    a.hostname AS agent_hostname

FROM audit_logs al

LEFT JOIN users u
ON al.actor_id = u.id
AND al.actor_type = 'user'

LEFT JOIN agents a
ON al.actor_id = a.id
AND al.actor_type = 'agent';

-- CLEANUP

CREATE OR REPLACE PROCEDURE cleanup_old_commands(
    days_to_keep INTEGER
)
LANGUAGE plpgsql
AS
$$
DECLARE
    deleted_rows INTEGER;
BEGIN

    DELETE FROM commands
    WHERE status = 'completed'
    AND completed_at < NOW() - (days_to_keep || ' days')::INTERVAL;

    GET DIAGNOSTICS deleted_rows = ROW_COUNT;

    INSERT INTO audit_logs (
        actor_type,
        action,
        target,
        details
    )
    VALUES (
        'system',
        'cleanup_completed_commands',
        'commands',
        jsonb_build_object(
            'days_to_keep', days_to_keep,
            'deleted_rows', deleted_rows
        )
    );

END;
$$;

-- SYSTEM STATS 

CREATE OR REPLACE FUNCTION get_system_stats()
RETURNS TABLE (
    total_users BIGINT,
    total_agents BIGINT,
    online_agents BIGINT,
    pending_commands BIGINT
)
LANGUAGE plpgsql
AS
$$
BEGIN

    RETURN QUERY
    SELECT
        (SELECT COUNT(*) FROM users WHERE is_active = TRUE),
        (SELECT COUNT(*) FROM agents),
        (SELECT COUNT(*) FROM agents WHERE status = 'online'),
        (SELECT COUNT(*) FROM commands WHERE status = 'pending');

END;
$$;