-- Comprehensive AI and user usage tracking system
-- Create AI usage tracking table
CREATE TABLE ai_usage (
    id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
    user_id UUID REFERENCES auth.users(id) ON DELETE CASCADE NOT NULL,
    book_id UUID REFERENCES books(id) ON DELETE SET NULL,
    chat_id UUID REFERENCES chats(id) ON DELETE SET NULL,
    
    -- Model information at time of usage
    model_name TEXT NOT NULL,
    model_tier TEXT NOT NULL CHECK (model_tier IN ('advanced', 'auto')),
    model_provider TEXT NOT NULL,
    
    -- Usage details
    prompt_tokens INTEGER DEFAULT 0 CHECK (prompt_tokens >= 0),
    completion_tokens INTEGER DEFAULT 0 CHECK (completion_tokens >= 0),
    total_tokens INTEGER DEFAULT 0 CHECK (total_tokens >= 0),
    
    -- Cost tracking (optional, in USD)
    cost_usd DECIMAL(10, 6) DEFAULT 0 CHECK (cost_usd >= 0),
    
    -- Request metadata
    request_type TEXT DEFAULT 'chat' CHECK (request_type IN ('chat', 'generation', 'editing', 'completion')),
    success BOOLEAN DEFAULT true NOT NULL,
    error_message TEXT,
    
    -- Timestamps
    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW() NOT NULL,
    billing_month DATE,
    
    -- Ensure total_tokens matches the sum
    CONSTRAINT check_total_tokens CHECK (total_tokens = prompt_tokens + completion_tokens)
);

-- Create user usage tracking table for plan limits
CREATE TABLE user_usage (
  id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
  user_id UUID REFERENCES auth.users(id) ON DELETE CASCADE NOT NULL,
  period_start TIMESTAMPTZ NOT NULL,
  period_end TIMESTAMPTZ NOT NULL,
  
  -- Usage counters
  advanced_prompts_used INTEGER DEFAULT 0 NOT NULL,
  auto_prompts_used INTEGER DEFAULT 0 NOT NULL,
  storage_used_gb DECIMAL(10,3) DEFAULT 0 NOT NULL,
  custom_templates_created INTEGER DEFAULT 0 NOT NULL,
  books_created INTEGER DEFAULT 0 NOT NULL,
  
  -- Metadata
  last_reset_at TIMESTAMPTZ DEFAULT NOW() NOT NULL,
  created_at TIMESTAMPTZ DEFAULT NOW() NOT NULL,
  updated_at TIMESTAMPTZ DEFAULT NOW() NOT NULL,
  
  -- Ensure one record per user per period
  UNIQUE(user_id, period_start, period_end)
);

-- Create usage logs table for detailed tracking
CREATE TABLE usage_logs (
  id UUID DEFAULT gen_random_uuid() PRIMARY KEY,
  user_id UUID REFERENCES auth.users(id) ON DELETE CASCADE NOT NULL,
  feature_type TEXT NOT NULL CHECK (feature_type IN (
    'advanced_prompts', 'auto_prompts', 'storage_gb', 'custom_templates', 'books'
  )),
  action TEXT NOT NULL CHECK (action IN ('increment', 'decrement', 'reset')),
  amount DECIMAL(10,3) NOT NULL,
  metadata JSONB DEFAULT '{}'::jsonb,
  created_at TIMESTAMPTZ DEFAULT NOW() NOT NULL
);

-- Create trigger function to set billing_month for ai_usage
CREATE OR REPLACE FUNCTION set_billing_month()
RETURNS TRIGGER AS $$
BEGIN
    NEW.billing_month := DATE_TRUNC('month', NEW.created_at)::DATE;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

-- Create trigger to automatically set billing_month
CREATE TRIGGER trigger_set_billing_month
    BEFORE INSERT ON ai_usage
    FOR EACH ROW
    EXECUTE FUNCTION set_billing_month();

-- Create trigger for updated_at on user_usage
CREATE TRIGGER update_user_usage_updated_at 
    BEFORE UPDATE ON user_usage 
    FOR EACH ROW 
    EXECUTE FUNCTION update_updated_at_column();

-- Create indexes for efficient querying
CREATE INDEX idx_ai_usage_user_id ON ai_usage(user_id);
CREATE INDEX idx_ai_usage_user_billing_month ON ai_usage(user_id, billing_month);
CREATE INDEX idx_ai_usage_model_tier ON ai_usage(user_id, model_tier, billing_month);
CREATE INDEX idx_ai_usage_created_at ON ai_usage(created_at DESC);
CREATE INDEX idx_ai_usage_book_id ON ai_usage(book_id) WHERE book_id IS NOT NULL;
CREATE INDEX idx_ai_usage_chat_id ON ai_usage(chat_id) WHERE chat_id IS NOT NULL;
CREATE INDEX idx_ai_usage_success ON ai_usage(user_id, success, billing_month);

-- User usage indexes
CREATE INDEX idx_user_usage_user_id ON user_usage(user_id);
CREATE INDEX idx_user_usage_period ON user_usage(user_id, period_start, period_end);
CREATE INDEX idx_user_usage_unique_period ON user_usage(user_id, period_start, period_end) WHERE period_start IS NOT NULL AND period_end IS NOT NULL;

-- Usage logs indexes
CREATE INDEX idx_usage_logs_user_id ON usage_logs(user_id);
CREATE INDEX idx_usage_logs_feature_type ON usage_logs(feature_type);
CREATE INDEX idx_usage_logs_created_at ON usage_logs(created_at);
CREATE INDEX idx_usage_logs_user_feature ON usage_logs(user_id, feature_type, created_at);

-- Add unique constraint to prevent duplicate tracking of the same request
-- (in case of race conditions or retry logic)
CREATE UNIQUE INDEX idx_ai_usage_dedup ON ai_usage(user_id, model_name, created_at, prompt_tokens, completion_tokens)
WHERE success = true;

-- Enable Row Level Security
ALTER TABLE ai_usage ENABLE ROW LEVEL SECURITY;
ALTER TABLE user_usage ENABLE ROW LEVEL SECURITY;
ALTER TABLE usage_logs ENABLE ROW LEVEL SECURITY;

-- AI Usage policies
CREATE POLICY "Users can view own AI usage" ON ai_usage
    FOR SELECT USING (auth.uid() = user_id);

CREATE POLICY "System can insert AI usage" ON ai_usage
    FOR INSERT WITH CHECK (auth.uid() = user_id);

CREATE POLICY "No updates allowed on AI usage" ON ai_usage FOR UPDATE USING (false);
CREATE POLICY "No deletes allowed on AI usage" ON ai_usage FOR DELETE USING (false);

-- User usage policies
CREATE POLICY "Users can view their own usage" ON user_usage
  FOR SELECT USING (auth.uid() = user_id);

CREATE POLICY "Users can insert their own usage" ON user_usage
  FOR INSERT WITH CHECK (auth.uid() = user_id);

CREATE POLICY "Users can update their own usage" ON user_usage
  FOR UPDATE USING (auth.uid() = user_id);

-- Usage logs policies
CREATE POLICY "Users can view their own usage logs" ON usage_logs
  FOR SELECT USING (auth.uid() = user_id);

CREATE POLICY "Users can insert their own usage logs" ON usage_logs
  FOR INSERT WITH CHECK (auth.uid() = user_id);

-- Function to get current month usage by tier
CREATE OR REPLACE FUNCTION get_monthly_ai_usage(p_user_id UUID, p_tier TEXT DEFAULT NULL)
RETURNS TABLE (
    tier TEXT,
    usage_count BIGINT,
    total_tokens BIGINT,
    total_cost DECIMAL
) AS $$
BEGIN
    RETURN QUERY
    SELECT 
        au.model_tier as tier,
        COUNT(*)::BIGINT as usage_count,
        COALESCE(SUM(au.total_tokens), 0)::BIGINT as total_tokens,
        COALESCE(SUM(au.cost_usd), 0)::DECIMAL as total_cost
    FROM ai_usage au
    WHERE au.user_id = p_user_id
        AND au.billing_month = DATE_TRUNC('month', NOW())
        AND au.success = true -- Only count successful requests
        AND (p_tier IS NULL OR au.model_tier = p_tier)
    GROUP BY au.model_tier;
END;
$$ LANGUAGE plpgsql;

-- Function to get usage stats for dashboard
CREATE OR REPLACE FUNCTION get_user_usage_stats(p_user_id UUID)
RETURNS JSON AS $$
DECLARE
    result JSON;
BEGIN
    SELECT json_build_object(
        'current_month', json_build_object(
            'advanced', COALESCE((
                SELECT json_build_object(
                    'count', usage_count,
                    'tokens', total_tokens,
                    'cost', total_cost
                )
                FROM get_monthly_ai_usage(p_user_id, 'advanced')
                WHERE tier = 'advanced'
            ), json_build_object('count', 0, 'tokens', 0, 'cost', 0)),
            'auto', COALESCE((
                SELECT json_build_object(
                    'count', usage_count,
                    'tokens', total_tokens,
                    'cost', total_cost
                )
                FROM get_monthly_ai_usage(p_user_id, 'auto')
                WHERE tier = 'auto'
            ), json_build_object('count', 0, 'tokens', 0, 'cost', 0))
        ),
        'total_books', (
            SELECT COUNT(*) FROM books WHERE user_id = p_user_id
        ),
        'total_words', (
            SELECT COALESCE(SUM(word_count), 0) FROM books WHERE user_id = p_user_id
        )
    ) INTO result;
    
    RETURN result;
END;
$$ LANGUAGE plpgsql;

-- Function to check if user can make a request (rate limiting)
CREATE OR REPLACE FUNCTION can_user_make_request(
    p_user_id UUID, 
    p_tier TEXT, 
    p_limit INTEGER
) RETURNS JSON AS $$
DECLARE
    current_usage BIGINT;
    result JSON;
BEGIN
    -- Get current month usage for the tier
    SELECT COALESCE(usage_count, 0) INTO current_usage
    FROM get_monthly_ai_usage(p_user_id, p_tier)
    WHERE tier = p_tier;
    
    -- Return usage info and whether request can proceed
    SELECT json_build_object(
        'can_proceed', (current_usage < p_limit OR p_limit = 0),
        'current_usage', current_usage,
        'limit', p_limit,
        'remaining', GREATEST(0, p_limit - current_usage)
    ) INTO result;
    
    RETURN result;
END;
$$ LANGUAGE plpgsql;

-- Note: Functions that depend on subscriptions table moved to 013_usage_functions_with_subscriptions.sql
-- This includes get_current_usage_period, get_or_create_current_usage, and increment_usage

-- Add helpful comments
COMMENT ON TABLE ai_usage IS 'Tracks detailed AI usage for billing and analytics';
COMMENT ON TABLE user_usage IS 'Tracks user usage against their pricing tier limits for current billing period';
COMMENT ON TABLE usage_logs IS 'Detailed log of all usage events for auditing and analytics';
COMMENT ON FUNCTION get_monthly_ai_usage(UUID, TEXT) IS 'Gets AI usage stats for current month by tier';
COMMENT ON FUNCTION get_user_usage_stats(UUID) IS 'Gets comprehensive usage stats for dashboard';
COMMENT ON FUNCTION can_user_make_request(UUID, TEXT, INTEGER) IS 'Checks if user can make AI request within limits';
-- Comments for moved functions are in 013_usage_functions_with_subscriptions.sql