-- ============================================================
-- NordaBiz - Migration 006: AI Usage Tracking
-- ============================================================
-- Created: 2026-01-11
-- Description:
--   Track AI (Gemini) API usage for monitoring and cost control
--   - Request logs with tokens and costs
--   - Daily/monthly aggregations
--   - Rate limit tracking
-- ============================================================

-- ============================================================
-- 1. AI USAGE LOGS - Individual API calls
-- ============================================================

CREATE TABLE IF NOT EXISTS ai_usage_logs (
    id SERIAL PRIMARY KEY,

    -- Request info
    request_type VARCHAR(50) NOT NULL,  -- 'chat', 'news_evaluation', 'user_creation', 'image_analysis'
    model VARCHAR(100) NOT NULL,        -- 'gemini-2.0-flash', 'gemini-1.5-pro', etc.

    -- Token counts
    tokens_input INTEGER DEFAULT 0,
    tokens_output INTEGER DEFAULT 0,
    tokens_total INTEGER GENERATED ALWAYS AS (tokens_input + tokens_output) STORED,

    -- Cost (in USD cents for precision)
    cost_cents DECIMAL(10, 4) DEFAULT 0,

    -- Context
    user_id INTEGER REFERENCES users(id),
    company_id INTEGER REFERENCES companies(id),
    related_entity_type VARCHAR(50),     -- 'zopk_news', 'chat_message', 'company', etc.
    related_entity_id INTEGER,

    -- Request details
    prompt_length INTEGER,
    response_length INTEGER,
    response_time_ms INTEGER,            -- How long the API call took

    -- Status
    success BOOLEAN DEFAULT TRUE,
    error_message TEXT,

    -- Timestamps
    created_at TIMESTAMP DEFAULT NOW()
);

-- Indexes for efficient querying
CREATE INDEX IF NOT EXISTS idx_ai_usage_logs_created_at ON ai_usage_logs(created_at);
CREATE INDEX IF NOT EXISTS idx_ai_usage_logs_request_type ON ai_usage_logs(request_type);
CREATE INDEX IF NOT EXISTS idx_ai_usage_logs_user_id ON ai_usage_logs(user_id);
CREATE INDEX IF NOT EXISTS idx_ai_usage_logs_date ON ai_usage_logs(DATE(created_at));

-- ============================================================
-- 2. AI USAGE DAILY SUMMARY - Pre-aggregated stats
-- ============================================================

CREATE TABLE IF NOT EXISTS ai_usage_daily (
    id SERIAL PRIMARY KEY,

    date DATE NOT NULL UNIQUE,

    -- Request counts by type
    chat_requests INTEGER DEFAULT 0,
    news_evaluation_requests INTEGER DEFAULT 0,
    user_creation_requests INTEGER DEFAULT 0,
    image_analysis_requests INTEGER DEFAULT 0,
    other_requests INTEGER DEFAULT 0,
    total_requests INTEGER DEFAULT 0,

    -- Token totals
    total_tokens_input INTEGER DEFAULT 0,
    total_tokens_output INTEGER DEFAULT 0,
    total_tokens INTEGER DEFAULT 0,

    -- Cost (in USD cents)
    total_cost_cents DECIMAL(10, 4) DEFAULT 0,

    -- Performance
    avg_response_time_ms INTEGER,
    error_count INTEGER DEFAULT 0,

    -- Timestamps
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW()
);

CREATE INDEX IF NOT EXISTS idx_ai_usage_daily_date ON ai_usage_daily(date);

-- ============================================================
-- 3. AI RATE LIMITS - Track quota usage
-- ============================================================

CREATE TABLE IF NOT EXISTS ai_rate_limits (
    id SERIAL PRIMARY KEY,

    -- Limit type
    limit_type VARCHAR(50) NOT NULL,     -- 'daily', 'hourly', 'per_minute'
    limit_scope VARCHAR(50) NOT NULL,    -- 'global', 'user', 'ip'
    scope_identifier VARCHAR(255),       -- user_id, ip address, or NULL for global

    -- Limits
    max_requests INTEGER NOT NULL,
    current_requests INTEGER DEFAULT 0,

    -- Reset
    reset_at TIMESTAMP NOT NULL,

    -- Timestamps
    created_at TIMESTAMP DEFAULT NOW(),
    updated_at TIMESTAMP DEFAULT NOW(),

    UNIQUE(limit_type, limit_scope, scope_identifier)
);

-- ============================================================
-- 4. GRANT PERMISSIONS
-- ============================================================

GRANT ALL ON TABLE ai_usage_logs TO nordabiz_app;
GRANT ALL ON TABLE ai_usage_daily TO nordabiz_app;
GRANT ALL ON TABLE ai_rate_limits TO nordabiz_app;
GRANT USAGE, SELECT ON SEQUENCE ai_usage_logs_id_seq TO nordabiz_app;
GRANT USAGE, SELECT ON SEQUENCE ai_usage_daily_id_seq TO nordabiz_app;
GRANT USAGE, SELECT ON SEQUENCE ai_rate_limits_id_seq TO nordabiz_app;

-- ============================================================
-- 5. HELPER FUNCTION - Update daily summary
-- ============================================================

CREATE OR REPLACE FUNCTION update_ai_usage_daily()
RETURNS TRIGGER AS $$
BEGIN
    INSERT INTO ai_usage_daily (
        date,
        chat_requests,
        news_evaluation_requests,
        user_creation_requests,
        image_analysis_requests,
        other_requests,
        total_requests,
        total_tokens_input,
        total_tokens_output,
        total_tokens,
        total_cost_cents,
        error_count,
        updated_at
    )
    VALUES (
        DATE(NEW.created_at),
        CASE WHEN NEW.request_type = 'chat' THEN 1 ELSE 0 END,
        CASE WHEN NEW.request_type = 'news_evaluation' THEN 1 ELSE 0 END,
        CASE WHEN NEW.request_type = 'user_creation' THEN 1 ELSE 0 END,
        CASE WHEN NEW.request_type = 'image_analysis' THEN 1 ELSE 0 END,
        CASE WHEN NEW.request_type NOT IN ('chat', 'news_evaluation', 'user_creation', 'image_analysis') THEN 1 ELSE 0 END,
        1,
        COALESCE(NEW.tokens_input, 0),
        COALESCE(NEW.tokens_output, 0),
        COALESCE(NEW.tokens_input, 0) + COALESCE(NEW.tokens_output, 0),
        COALESCE(NEW.cost_cents, 0),
        CASE WHEN NEW.success = FALSE THEN 1 ELSE 0 END,
        NOW()
    )
    ON CONFLICT (date) DO UPDATE SET
        chat_requests = ai_usage_daily.chat_requests + CASE WHEN NEW.request_type = 'chat' THEN 1 ELSE 0 END,
        news_evaluation_requests = ai_usage_daily.news_evaluation_requests + CASE WHEN NEW.request_type = 'news_evaluation' THEN 1 ELSE 0 END,
        user_creation_requests = ai_usage_daily.user_creation_requests + CASE WHEN NEW.request_type = 'user_creation' THEN 1 ELSE 0 END,
        image_analysis_requests = ai_usage_daily.image_analysis_requests + CASE WHEN NEW.request_type = 'image_analysis' THEN 1 ELSE 0 END,
        other_requests = ai_usage_daily.other_requests + CASE WHEN NEW.request_type NOT IN ('chat', 'news_evaluation', 'user_creation', 'image_analysis') THEN 1 ELSE 0 END,
        total_requests = ai_usage_daily.total_requests + 1,
        total_tokens_input = ai_usage_daily.total_tokens_input + COALESCE(NEW.tokens_input, 0),
        total_tokens_output = ai_usage_daily.total_tokens_output + COALESCE(NEW.tokens_output, 0),
        total_tokens = ai_usage_daily.total_tokens + COALESCE(NEW.tokens_input, 0) + COALESCE(NEW.tokens_output, 0),
        total_cost_cents = ai_usage_daily.total_cost_cents + COALESCE(NEW.cost_cents, 0),
        error_count = ai_usage_daily.error_count + CASE WHEN NEW.success = FALSE THEN 1 ELSE 0 END,
        updated_at = NOW();

    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- Create trigger to auto-update daily summary
DROP TRIGGER IF EXISTS trigger_ai_usage_daily ON ai_usage_logs;
CREATE TRIGGER trigger_ai_usage_daily
    AFTER INSERT ON ai_usage_logs
    FOR EACH ROW
    EXECUTE FUNCTION update_ai_usage_daily();
