CREATE EXTENSION IF NOT EXISTS "uuid-ossp";
CREATE EXTENSION IF NOT EXISTS "pg_trgm";

CREATE TABLE IF NOT EXISTS sessions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    tenant_id VARCHAR(64) DEFAULT 'default',
    created_at TIMESTAMPTZ DEFAULT NOW(),
    last_active_at TIMESTAMPTZ DEFAULT NOW(),
    expires_at TIMESTAMPTZ DEFAULT NOW() + INTERVAL '7 days',
    budget_usd NUMERIC(10, 6) DEFAULT 10.0,
    spent_usd NUMERIC(10, 6) DEFAULT 0.0,
    status VARCHAR(32) DEFAULT 'active',
    metadata JSONB DEFAULT '{}',
    context_json JSONB DEFAULT '[]',
    model_name VARCHAR(128),
    reasoning_effort VARCHAR(32),
    max_iterations INTEGER DEFAULT 300,
    yolo_mode BOOLEAN DEFAULT FALSE,
    CONSTRAINT positive_budget CHECK (budget_usd >= 0),
    CONSTRAINT positive_spent CHECK (spent_usd >= 0)
);

CREATE TABLE IF NOT EXISTS requests (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    session_id UUID REFERENCES sessions(id) ON DELETE CASCADE,
    tenant_id VARCHAR(64) DEFAULT 'default',
    trace_id VARCHAR(128),
    model VARCHAR(128) NOT NULL,
    provider VARCHAR(64) NOT NULL,
    operation_type VARCHAR(32) DEFAULT 'chat',
    input_tokens INTEGER DEFAULT 0,
    output_tokens INTEGER DEFAULT 0,
    total_tokens INTEGER GENERATED ALWAYS AS (input_tokens + output_tokens) STORED,
    cost_usd NUMERIC(10, 6) DEFAULT 0.0,
    cached BOOLEAN DEFAULT FALSE,
    latency_ms INTEGER,
    queue_time_ms INTEGER DEFAULT 0,
    request_payload JSONB,
    response_payload JSONB,
    tool_calls JSONB DEFAULT '[]',
    finish_reason VARCHAR(32),
    rate_limited BOOLEAN DEFAULT FALSE,
    circuit_open BOOLEAN DEFAULT FALSE,
    retry_count INTEGER DEFAULT 0,
    created_at TIMESTAMPTZ DEFAULT NOW(),
    CONSTRAINT non_negative_tokens CHECK (input_tokens >= 0 AND output_tokens >= 0)
);

CREATE TABLE IF NOT EXISTS circuit_events (
    id SERIAL PRIMARY KEY,
    provider VARCHAR(64) NOT NULL,
    event_type VARCHAR(32) NOT NULL,
    failure_count INTEGER DEFAULT 0,
    details JSONB DEFAULT '{}',
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS budget_events (
    id SERIAL PRIMARY KEY,
    session_id UUID REFERENCES sessions(id) ON DELETE CASCADE,
    event_type VARCHAR(32) NOT NULL,
    threshold_percent INTEGER,
    spent_before NUMERIC(10, 6),
    spent_after NUMERIC(10, 6),
    budget NUMERIC(10, 6),
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS tool_executions (
    id UUID PRIMARY KEY DEFAULT uuid_generate_v4(),
    request_id UUID REFERENCES requests(id) ON DELETE CASCADE,
    tool_name VARCHAR(128) NOT NULL,
    tool_args JSONB,
    result_json JSONB,
    success BOOLEAN DEFAULT TRUE,
    error_message TEXT,
    latency_ms INTEGER,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE TABLE IF NOT EXISTS cache_events (
    id SERIAL PRIMARY KEY,
    cache_key VARCHAR(256) NOT NULL,
    cache_type VARCHAR(32) NOT NULL,
    hit BOOLEAN NOT NULL,
    ttl_seconds INTEGER,
    size_bytes INTEGER,
    created_at TIMESTAMPTZ DEFAULT NOW()
);

CREATE INDEX IF NOT EXISTS idx_sessions_tenant ON sessions(tenant_id);
CREATE INDEX IF NOT EXISTS idx_sessions_status ON sessions(status);
CREATE INDEX IF NOT EXISTS idx_sessions_expires ON sessions(expires_at);
CREATE INDEX IF NOT EXISTS idx_sessions_created ON sessions(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_sessions_active ON sessions(tenant_id, status) WHERE status = 'active';
CREATE INDEX IF NOT EXISTS idx_sessions_metadata ON sessions USING GIN (metadata);

CREATE INDEX IF NOT EXISTS idx_requests_session ON requests(session_id);
CREATE INDEX IF NOT EXISTS idx_requests_tenant ON requests(tenant_id);
CREATE INDEX IF NOT EXISTS idx_requests_trace ON requests(trace_id);
CREATE INDEX IF NOT EXISTS idx_requests_provider ON requests(provider);
CREATE INDEX IF NOT EXISTS idx_requests_created ON requests(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_requests_model ON requests(model);
CREATE INDEX IF NOT EXISTS idx_requests_payload ON requests USING GIN (request_payload);

CREATE INDEX IF NOT EXISTS idx_circuit_provider ON circuit_events(provider);
CREATE INDEX IF NOT EXISTS idx_circuit_created ON circuit_events(created_at DESC);
CREATE INDEX IF NOT EXISTS idx_budget_session ON budget_events(session_id);
CREATE INDEX IF NOT EXISTS idx_tool_request ON tool_executions(request_id);
CREATE INDEX IF NOT EXISTS idx_tool_name ON tool_executions(tool_name);
CREATE INDEX IF NOT EXISTS idx_cache_key ON cache_events(cache_key);
CREATE INDEX IF NOT EXISTS idx_cache_created ON cache_events(created_at DESC);

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

DROP TRIGGER IF EXISTS trigger_update_last_active ON sessions;
CREATE TRIGGER trigger_update_last_active
    BEFORE UPDATE ON sessions
    FOR EACH ROW
    EXECUTE FUNCTION update_last_active();

CREATE MATERIALIZED VIEW IF NOT EXISTS mv_daily_costs AS
SELECT 
    tenant_id,
    DATE(created_at) as date,
    provider,
    model,
    COUNT(*) as request_count,
    SUM(input_tokens) as total_input_tokens,
    SUM(output_tokens) as total_output_tokens,
    SUM(cost_usd) as total_cost_usd,
    AVG(latency_ms)::NUMERIC(10,2) as avg_latency_ms,
    SUM(CASE WHEN cached THEN 1 ELSE 0 END) as cache_hits,
    SUM(CASE WHEN NOT cached THEN 1 ELSE 0 END) as cache_misses
FROM requests
GROUP BY tenant_id, DATE(created_at), provider, model;

CREATE UNIQUE INDEX IF NOT EXISTS idx_mv_daily_costs ON mv_daily_costs(tenant_id, date, provider, model);

CREATE OR REPLACE FUNCTION refresh_daily_costs()
RETURNS void AS $$
BEGIN
    REFRESH MATERIALIZED VIEW CONCURRENTLY mv_daily_costs;
END;
$$ LANGUAGE plpgsql;

ALTER TABLE sessions ENABLE ROW LEVEL SECURITY;
ALTER TABLE requests ENABLE ROW LEVEL SECURITY;

CREATE POLICY tenant_isolation_sessions ON sessions
    USING (tenant_id = current_setting('app.current_tenant', true));
CREATE POLICY tenant_isolation_requests ON requests
    USING (tenant_id = current_setting('app.current_tenant', true));

GRANT ALL PRIVILEGES ON ALL TABLES IN SCHEMA public TO ml_intern;
GRANT ALL PRIVILEGES ON ALL SEQUENCES IN SCHEMA public TO ml_intern;