-- Dedicated junction table for trace-entity many-to-many relationships
-- Partitioned by tagged_at for efficient time-based pruning

CREATE TABLE IF NOT EXISTS scouter.trace_entities (
    trace_id BYTEA NOT NULL CHECK (octet_length(trace_id) = 16),
    entity_uid BYTEA NOT NULL CHECK (octet_length(entity_uid) = 16),
    tagged_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    PRIMARY KEY (entity_uid, tagged_at, trace_id)
) PARTITION BY RANGE (tagged_at);

-- Create partitioned table using pg_partman (1-day partitions)
SELECT scouter.create_parent(
    'scouter.trace_entities',
    'tagged_at',
    '1 day'
);

-- Configure partition maintenance: premake 7 days ahead, retain 30 days
UPDATE scouter.part_config
    SET premake = 7,
    retention = '30 days',
    retention_keep_table = FALSE
WHERE parent_table = 'scouter.trace_entities';

-- Hot path index: lookup trace_ids by entity_uid
-- This index is automatically created on each partition by pg_partman
-- Ordering by tagged_at DESC supports time-bounded queries efficiently
CREATE INDEX idx_trace_entities_entity_lookup
ON scouter.trace_entities (entity_uid, tagged_at DESC)
INCLUDE (trace_id);

-- Secondary index: reverse lookup (trace -> entities), useful for some analytics
CREATE INDEX idx_trace_entities_trace_lookup
ON scouter.trace_entities (trace_id, tagged_at DESC)
INCLUDE (entity_uid);

-- Comments for context
COMMENT ON TABLE scouter.trace_entities IS
'Many-to-many junction table mapping traces to entities (prompts, data sources, services, etc.). Optimized for high-volume queries.';


-- Function: Get traces for a single entity (hot path)
CREATE OR REPLACE FUNCTION scouter.get_traces_by_entity(
    p_entity_uid BYTEA, -- Expecting 16-byte binary UUID
    p_start_time TIMESTAMPTZ DEFAULT NOW() - INTERVAL '24 hours',
    p_end_time TIMESTAMPTZ DEFAULT NOW(),
    p_limit INTEGER DEFAULT 50,
    p_cursor_start_time TIMESTAMPTZ DEFAULT NULL,
    p_cursor_trace_id BYTEA DEFAULT NULL
)
RETURNS TABLE (
    trace_id TEXT,
    service_name TEXT,
    scope_name TEXT,
    scope_version TEXT,
    root_operation TEXT,
    start_time TIMESTAMPTZ,
    end_time TIMESTAMPTZ,
    duration_ms BIGINT,
    status_code INTEGER,
    status_message TEXT,
    span_count BIGINT,
    has_errors BOOLEAN,
    error_count BIGINT,
    resource_attributes JSONB
)
LANGUAGE SQL
STABLE
AS $$

    WITH entity_traces AS (
        SELECT DISTINCT te.trace_id
        FROM scouter.trace_entities te
        WHERE te.entity_uid = p_entity_uid
          AND te.tagged_at >= p_start_time
          AND te.tagged_at <= p_end_time
    )
    SELECT
        encode(t.trace_id, 'hex') as trace_id,
        t.service_name,
        t.scope_name,
        t.scope_version,
        t.root_operation,
        t.start_time,
        t.end_time,
        t.duration_ms,
        t.status_code,
        t.status_message,
        t.span_count,
        t.error_count > 0 AS has_errors,
        t.error_count,
        t.resource_attributes
    FROM scouter.traces t
    WHERE t.trace_id IN (SELECT trace_id FROM entity_traces)
      AND t.bucket_time >= p_start_time
      AND t.bucket_time <= p_end_time
      AND (
          p_cursor_start_time IS NULL OR
          (t.start_time, t.trace_id) < (p_cursor_start_time, p_cursor_trace_id)
      )
    ORDER BY t.start_time DESC, t.trace_id
    LIMIT p_limit;
$$;


-- Function: Get traces matching ALL entities (service view with multiple filters)
CREATE OR REPLACE FUNCTION scouter.get_traces_by_multiple_entities(
    p_entity_uids BYTEA[],
    p_match_all BOOLEAN DEFAULT TRUE,
    p_start_time TIMESTAMPTZ DEFAULT NOW() - INTERVAL '24 hours',
    p_end_time TIMESTAMPTZ DEFAULT NOW(),
    p_limit INTEGER DEFAULT 50
)
RETURNS TABLE (
    trace_id TEXT,
    service_name TEXT,
    root_operation TEXT,
    start_time TIMESTAMPTZ,
    duration_ms BIGINT,
    error_count BIGINT,
    matched_entities TEXT[]
)
LANGUAGE SQL
STABLE
AS $$

    WITH entity_uid_bytes AS (
        SELECT uid as entity_uid_byte, uid
        FROM unnest(p_entity_uids) AS uid
    ),
    entity_traces AS (
        SELECT
            te.trace_id,
            array_agg(DISTINCT eub.uid) as matched_entities,
            COUNT(DISTINCT te.entity_uid) as match_count
        FROM scouter.trace_entities te
        INNER JOIN entity_uid_bytes eub ON te.entity_uid = eub.entity_uid_byte
        WHERE te.tagged_at >= p_start_time
          AND te.tagged_at <= p_end_time
        GROUP BY te.trace_id
        HAVING
            CASE
                WHEN p_match_all THEN
                    COUNT(DISTINCT te.entity_uid) = array_length(p_entity_uids, 1)
                ELSE
                    COUNT(DISTINCT te.entity_uid) > 0
            END
    )
    SELECT
        encode(t.trace_id, 'hex') as trace_id,
        t.service_name,
        t.root_operation,
        t.start_time,
        t.duration_ms,
        t.error_count,
        et.matched_entities
    FROM scouter.traces t
    INNER JOIN entity_traces et ON t.trace_id = et.trace_id
    WHERE t.bucket_time >= p_start_time
      AND t.bucket_time <= p_end_time
    ORDER BY t.start_time DESC
    LIMIT p_limit;
$$;

CREATE OR REPLACE FUNCTION scouter.get_traces_paginated(
    p_service_name TEXT DEFAULT NULL,
    p_has_errors BOOLEAN DEFAULT NULL,
    p_start_time TIMESTAMPTZ DEFAULT NULL,
    p_end_time TIMESTAMPTZ DEFAULT NULL,
    p_limit INTEGER DEFAULT 50,
    p_cursor_start_time TIMESTAMPTZ DEFAULT NULL,
    p_cursor_trace_id BYTEA DEFAULT NULL,
    p_direction TEXT DEFAULT 'next',
    p_attribute_filters JSONB DEFAULT NULL,
    p_trace_ids BYTEA[] DEFAULT NULL,
    p_match_all_attributes BOOLEAN DEFAULT FALSE,
    p_entity_uid BYTEA DEFAULT NULL
)
RETURNS TABLE (
    trace_id TEXT,
    service_name TEXT,
    scope_name TEXT,
    scope_version TEXT,
    root_operation TEXT,
    start_time TIMESTAMPTZ,
    end_time TIMESTAMPTZ,
    duration_ms BIGINT,
    status_code INTEGER,
    status_message TEXT,
    span_count BIGINT,
    has_errors BOOLEAN,
    error_count BIGINT,
    resource_attributes JSONB
)
LANGUAGE SQL
STABLE
AS $$
    -- Entity filtering (when needed)
    WITH entity_filter AS (
        SELECT te.trace_id
        FROM scouter.trace_entities te
        WHERE te.entity_uid = p_entity_uid
          AND te.tagged_at >= COALESCE(p_start_time, NOW() - INTERVAL '24 hours')
          AND te.tagged_at <= COALESCE(p_end_time, NOW())
    ),
    -- Primary filtering on traces table (fast!)
    base_traces AS (
        SELECT
            t.trace_id,
            t.service_name,
            t.scope_name,
            t.scope_version,
            t.root_operation,
            t.start_time,
            t.end_time,
            t.duration_ms,
            t.status_code,
            t.status_message,
            t.span_count,
            t.error_count > 0 AS has_errors,
            t.error_count,
            t.resource_attributes
        FROM scouter.traces t
        WHERE
            t.bucket_time >= COALESCE(p_start_time, NOW() - INTERVAL '24 hours')
            AND t.bucket_time <= COALESCE(p_end_time, NOW())
            AND (p_service_name IS NULL OR t.service_name = p_service_name)
            AND (p_has_errors IS NULL OR (p_has_errors AND t.error_count > 0) OR (NOT p_has_errors AND t.error_count = 0))
            AND (p_trace_ids IS NULL OR t.trace_id = ANY(p_trace_ids))
            AND (p_entity_uid IS NULL OR t.trace_id IN (SELECT trace_id FROM entity_filter))
            AND (
                CASE
                    WHEN p_direction = 'next' THEN
                        (p_cursor_start_time IS NULL) OR
                        (t.start_time, t.trace_id) < (p_cursor_start_time, p_cursor_trace_id)
                    WHEN p_direction = 'previous' THEN
                        (p_cursor_start_time IS NULL) OR
                        (t.start_time, t.trace_id) > (p_cursor_start_time, p_cursor_trace_id)
                    ELSE TRUE
                END
            )
        ORDER BY
            CASE WHEN p_direction = 'previous' THEN t.start_time END ASC,
            CASE WHEN p_direction = 'next' THEN t.start_time END DESC,
            t.trace_id
        LIMIT p_limit + 1
    ),

    filtered_traces AS (
        SELECT bt.*
        FROM base_traces bt
        WHERE
            p_attribute_filters IS NULL
            OR EXISTS (
                SELECT 1
                FROM scouter.spans s
                WHERE s.trace_id = bt.trace_id
                  AND s.start_time >= bt.start_time - INTERVAL '1 hour'
                  AND s.start_time <= bt.start_time + INTERVAL '1 hour'
                  AND scouter.match_span_attributes(s.attributes, p_attribute_filters, p_match_all_attributes)
                LIMIT 1
            )
    )

    SELECT
        encode(ft.trace_id, 'hex') as trace_id,
        ft.service_name,
        ft.scope_name,
        ft.scope_version,
        ft.root_operation,
        ft.start_time,
        ft.end_time,
        ft.duration_ms,
        ft.status_code,
        ft.status_message,
        ft.span_count,
        ft.has_errors,
        ft.error_count,
        ft.resource_attributes
    FROM filtered_traces ft;
$$;
