database/schema.py

"""
Database schema definitions for FleetMind
Contains PostgreSQL schema for all database tables
"""

SCHEMA_SQL = """
-- ============================================
-- ORDERS TABLE
-- ============================================

CREATE TABLE IF NOT EXISTS orders (
    order_id VARCHAR(50) PRIMARY KEY,
    customer_name VARCHAR(255) NOT NULL,
    customer_phone VARCHAR(20),
    customer_email VARCHAR(255),
    pickup_address TEXT,
    pickup_lat DECIMAL(10, 8),
    pickup_lng DECIMAL(11, 8),
    delivery_address TEXT NOT NULL,
    delivery_lat DECIMAL(10, 8),
    delivery_lng DECIMAL(11, 8),
    time_window_start TIMESTAMP,
    time_window_end TIMESTAMP,
    priority VARCHAR(20) CHECK(priority IN ('standard', 'express', 'urgent')) DEFAULT 'standard',
    weight_kg DECIMAL(10, 2),
    volume_m3 DECIMAL(10, 3),
    special_instructions TEXT,
    status VARCHAR(20) CHECK(status IN ('pending', 'assigned', 'in_transit', 'delivered', 'failed', 'cancelled')) DEFAULT 'pending',
    assigned_driver_id VARCHAR(50),
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    delivered_at TIMESTAMP,

    -- Additional fields for order management
    order_value DECIMAL(10, 2),
    payment_status VARCHAR(20) CHECK(payment_status IN ('pending', 'paid', 'cod')) DEFAULT 'pending',
    requires_signature BOOLEAN DEFAULT FALSE,
    is_fragile BOOLEAN DEFAULT FALSE,
    requires_cold_storage BOOLEAN DEFAULT FALSE
);

-- ============================================
-- INDEXES FOR ORDERS TABLE
-- ============================================

CREATE INDEX IF NOT EXISTS idx_orders_status ON orders(status);
CREATE INDEX IF NOT EXISTS idx_orders_priority ON orders(priority);
CREATE INDEX IF NOT EXISTS idx_orders_assigned_driver ON orders(assigned_driver_id);
CREATE INDEX IF NOT EXISTS idx_orders_created_at ON orders(created_at);
CREATE INDEX IF NOT EXISTS idx_orders_time_window ON orders(time_window_start, time_window_end);

-- ============================================
-- TRIGGER TO UPDATE updated_at TIMESTAMP
-- ============================================

CREATE OR REPLACE FUNCTION update_updated_at_column()
RETURNS TRIGGER AS $$
BEGIN
    NEW.updated_at = CURRENT_TIMESTAMP;
    RETURN NEW;
END;
$$ language 'plpgsql';

DROP TRIGGER IF EXISTS update_orders_timestamp ON orders;
CREATE TRIGGER update_orders_timestamp
    BEFORE UPDATE ON orders
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();
"""

# Additional schema components can be added here as needed
DRIVERS_SCHEMA = """
-- ============================================
-- DRIVERS TABLE
-- ============================================

CREATE TABLE IF NOT EXISTS drivers (
    driver_id VARCHAR(50) PRIMARY KEY,
    name VARCHAR(255) NOT NULL,
    phone VARCHAR(20),
    email VARCHAR(255),
    current_lat DECIMAL(10, 8),
    current_lng DECIMAL(11, 8),
    last_location_update TIMESTAMP,
    status VARCHAR(20) CHECK(status IN ('active', 'busy', 'offline', 'unavailable')) DEFAULT 'active',
    vehicle_type VARCHAR(50),
    vehicle_plate VARCHAR(20),
    capacity_kg DECIMAL(10, 2),
    capacity_m3 DECIMAL(10, 3),
    skills JSONB,  -- JSON array: ["medical_certified", "refrigerated"]
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

CREATE INDEX IF NOT EXISTS idx_drivers_status ON drivers(status);

DROP TRIGGER IF EXISTS update_drivers_timestamp ON drivers;
CREATE TRIGGER update_drivers_timestamp
    BEFORE UPDATE ON drivers
    FOR EACH ROW
    EXECUTE FUNCTION update_updated_at_column();
"""

ASSIGNMENTS_SCHEMA = """
-- ============================================
-- ASSIGNMENTS TABLE
-- ============================================

CREATE TABLE IF NOT EXISTS assignments (
    assignment_id VARCHAR(50) PRIMARY KEY,
    order_id VARCHAR(50) NOT NULL,
    driver_id VARCHAR(50) NOT NULL,
    assigned_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    sequence_number INTEGER,  -- Order in driver's route
    estimated_arrival TIMESTAMP,
    actual_arrival TIMESTAMP,
    estimated_distance_km DECIMAL(10, 2),
    actual_distance_km DECIMAL(10, 2),
    status VARCHAR(20) CHECK(status IN ('assigned', 'in_progress', 'completed', 'failed')) DEFAULT 'assigned',
    notes TEXT,
    FOREIGN KEY (order_id) REFERENCES orders(order_id) ON DELETE CASCADE,
    FOREIGN KEY (driver_id) REFERENCES drivers(driver_id) ON DELETE CASCADE
);

CREATE INDEX IF NOT EXISTS idx_assignments_driver ON assignments(driver_id);
CREATE INDEX IF NOT EXISTS idx_assignments_order ON assignments(order_id);
CREATE INDEX IF NOT EXISTS idx_assignments_status ON assignments(status);
"""

EXCEPTIONS_SCHEMA = """
-- ============================================
-- EXCEPTIONS TABLE
-- ============================================

CREATE TABLE IF NOT EXISTS exceptions (
    exception_id VARCHAR(50) PRIMARY KEY,
    exception_type VARCHAR(50) CHECK(exception_type IN ('driver_breakdown', 'traffic_delay', 'weather', 'customer_unavailable', 'other')),
    severity VARCHAR(20) CHECK(severity IN ('low', 'medium', 'high', 'critical')),
    description TEXT,
    affected_orders JSONB,  -- JSON array of order IDs
    affected_drivers JSONB,  -- JSON array of driver IDs
    detected_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    resolved_at TIMESTAMP,
    resolution_method JSONB,  -- JSON: {agent: "...", actions: [...]}
    coordinator_override BOOLEAN DEFAULT FALSE,
    resolution_time_seconds INTEGER
);

CREATE INDEX IF NOT EXISTS idx_exceptions_type ON exceptions(exception_type);
CREATE INDEX IF NOT EXISTS idx_exceptions_severity ON exceptions(severity);
CREATE INDEX IF NOT EXISTS idx_exceptions_detected_at ON exceptions(detected_at);
"""

AGENT_DECISIONS_SCHEMA = """
-- ============================================
-- AGENT DECISIONS TABLE (For Agent Learning)
-- ============================================

CREATE TABLE IF NOT EXISTS agent_decisions (
    decision_id VARCHAR(50) PRIMARY KEY,
    session_id VARCHAR(50),
    task_type VARCHAR(50),
    situation_context JSONB,  -- JSON
    options_generated JSONB,  -- JSON array
    selected_option JSONB,  -- JSON
    reasoning TEXT,
    approval_status VARCHAR(20) CHECK(approval_status IN ('approved', 'rejected', 'modified')),
    coordinator_feedback TEXT,
    outcome_success BOOLEAN,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
);

CREATE INDEX IF NOT EXISTS idx_decisions_session ON agent_decisions(session_id);
CREATE INDEX IF NOT EXISTS idx_decisions_task_type ON agent_decisions(task_type);
"""

METRICS_SCHEMA = """
-- ============================================
-- PERFORMANCE METRICS TABLE
-- ============================================

CREATE TABLE IF NOT EXISTS metrics (
    metric_id VARCHAR(50) PRIMARY KEY,
    date DATE,
    total_orders INTEGER,
    on_time_deliveries INTEGER,
    failed_deliveries INTEGER,
    exceptions_count INTEGER,
    avg_resolution_time_seconds DECIMAL(10, 2),
    coordinator_time_saved_minutes INTEGER,
    fuel_cost_savings DECIMAL(10, 2)
);

CREATE INDEX IF NOT EXISTS idx_metrics_date ON metrics(date);
"""


def get_full_schema() -> str:
    """
    Get the complete database schema including all tables.

    Returns:
        str: Complete SQL schema
    """
    return (
        SCHEMA_SQL + "\n" +
        DRIVERS_SCHEMA + "\n" +
        ASSIGNMENTS_SCHEMA + "\n" +
        EXCEPTIONS_SCHEMA + "\n" +
        AGENT_DECISIONS_SCHEMA + "\n" +
        METRICS_SCHEMA
    )