-- ============================================
-- AI Usage Tracking Table
-- Records each AI API call for billing/analytics
-- ============================================

CREATE TABLE ai.usage (
    id UUID PRIMARY KEY DEFAULT gen_random_uuid(),
    tenant_id UUID NOT NULL REFERENCES tenants.tenants(id) ON DELETE CASCADE,
    user_id UUID NOT NULL REFERENCES users.users(id) ON DELETE CASCADE,

    -- Request details
    provider ai.ai_provider NOT NULL,
    model VARCHAR(100) NOT NULL,
    model_type ai.ai_model_type NOT NULL DEFAULT 'chat',
    status ai.usage_status NOT NULL DEFAULT 'pending',

    -- Token counts
    input_tokens INTEGER NOT NULL DEFAULT 0,
    output_tokens INTEGER NOT NULL DEFAULT 0,
    total_tokens INTEGER GENERATED ALWAYS AS (input_tokens + output_tokens) STORED,

    -- Cost tracking (in USD, 6 decimal precision)
    cost_input NUMERIC(12,6) NOT NULL DEFAULT 0,
    cost_output NUMERIC(12,6) NOT NULL DEFAULT 0,
    cost_total NUMERIC(12,6) GENERATED ALWAYS AS (cost_input + cost_output) STORED,

    -- Performance metrics
    latency_ms INTEGER,
    started_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
    completed_at TIMESTAMPTZ,

    -- Request metadata
    request_id VARCHAR(100),
    endpoint VARCHAR(50),
    error_message TEXT,

    -- Additional context
    metadata JSONB DEFAULT '{}',

    -- Timestamps
    created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);

-- Indexes for common queries
CREATE INDEX idx_ai_usage_tenant ON ai.usage(tenant_id);
CREATE INDEX idx_ai_usage_user ON ai.usage(user_id);
CREATE INDEX idx_ai_usage_tenant_created ON ai.usage(tenant_id, created_at DESC);
CREATE INDEX idx_ai_usage_model ON ai.usage(model);
CREATE INDEX idx_ai_usage_status ON ai.usage(status);
CREATE INDEX idx_ai_usage_created_at ON ai.usage(created_at DESC);

-- Index for monthly queries (used by get_current_month_usage function)
-- Note: Partial index with CURRENT_DATE removed as it's not IMMUTABLE
CREATE INDEX idx_ai_usage_monthly ON ai.usage(tenant_id, created_at, status)
    WHERE status = 'completed';

-- RLS
ALTER TABLE ai.usage ENABLE ROW LEVEL SECURITY;

CREATE POLICY ai_usage_tenant_isolation ON ai.usage
    FOR ALL
    USING (tenant_id = current_setting('app.current_tenant_id', true)::UUID);

-- Comments
COMMENT ON TABLE ai.usage IS 'AI API usage tracking for billing and analytics';
COMMENT ON COLUMN ai.usage.input_tokens IS 'Number of input/prompt tokens';
COMMENT ON COLUMN ai.usage.output_tokens IS 'Number of output/completion tokens';
COMMENT ON COLUMN ai.usage.cost_input IS 'Cost for input tokens in USD';
COMMENT ON COLUMN ai.usage.cost_output IS 'Cost for output tokens in USD';
COMMENT ON COLUMN ai.usage.latency_ms IS 'Request latency in milliseconds';


-- ============================================
-- Monthly Usage Summary View
-- For efficient billing queries
-- ============================================

CREATE VIEW ai.monthly_usage AS
SELECT
    tenant_id,
    date_trunc('month', created_at) AS month,
    COUNT(*) AS request_count,
    SUM(input_tokens) AS total_input_tokens,
    SUM(output_tokens) AS total_output_tokens,
    SUM(input_tokens + output_tokens) AS total_tokens,
    SUM(cost_input + cost_output) AS total_cost,
    AVG(latency_ms) AS avg_latency_ms
FROM ai.usage
WHERE status = 'completed'
GROUP BY tenant_id, date_trunc('month', created_at);

COMMENT ON VIEW ai.monthly_usage IS 'Monthly AI usage aggregation per tenant';


-- ============================================
-- Function: Get tenant usage for current month
-- ============================================

CREATE OR REPLACE FUNCTION ai.get_current_month_usage(p_tenant_id UUID)
RETURNS TABLE (
    request_count BIGINT,
    total_input_tokens BIGINT,
    total_output_tokens BIGINT,
    total_tokens BIGINT,
    total_cost NUMERIC,
    avg_latency_ms NUMERIC
) AS $$
BEGIN
    RETURN QUERY
    SELECT
        COUNT(*)::BIGINT,
        COALESCE(SUM(input_tokens), 0)::BIGINT,
        COALESCE(SUM(output_tokens), 0)::BIGINT,
        COALESCE(SUM(input_tokens + output_tokens), 0)::BIGINT,
        COALESCE(SUM(cost_input + cost_output), 0)::NUMERIC,
        COALESCE(AVG(latency_ms), 0)::NUMERIC
    FROM ai.usage
    WHERE tenant_id = p_tenant_id
      AND status = 'completed'
      AND created_at >= date_trunc('month', CURRENT_DATE);
END;
$$ LANGUAGE plpgsql STABLE;

COMMENT ON FUNCTION ai.get_current_month_usage IS 'Get AI usage for current month for a tenant';