import duckdb
import os
from pathlib import Path
from datetime import datetime
import uuid

# ─── Tenant Database Manager ───
class TenantDatabaseManager:
    """Multi-tenant database manager: each tenant gets an independent DuckDB file"""
    
    def __init__(self, data_dir: str = "/data/tenants"):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        # Global metadata database
        self.meta_conn = duckdb.connect(str(self.data_dir / "_meta.duckdb"))
        self._init_meta()
    
    def _init_meta(self):
        """Initialize tenant metadata table"""
        self.meta_conn.execute("""
            CREATE TABLE IF NOT EXISTS tenants (
                tenant_id     VARCHAR PRIMARY KEY,
                tenant_name   VARCHAR NOT NULL,
                plan          VARCHAR DEFAULT 'free',
                created_at    TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                db_path       VARCHAR NOT NULL,
                status        VARCHAR DEFAULT 'active',
                data_size_mb  DOUBLE DEFAULT 0,
                max_memory_mb INTEGER DEFAULT 512
            )
        """)
    
    def create_tenant(self, tenant_name: str, plan: str = "free") -> str:
        """Create a new tenant: register + initialize database"""
        tenant_id = f"t_{uuid.uuid4().hex[:12]}"
        db_path = str(self.data_dir / f"{tenant_id}.duckdb")
        
        # Register tenant
        self.meta_conn.execute("""
            INSERT INTO tenants (tenant_id, tenant_name, plan, db_path)
            VALUES (?, ?, ?, ?)
        """, [tenant_id, tenant_name, plan, db_path])
        
        # Initialize tenant database
        self._init_tenant_db(db_path, plan)
        
        return tenant_id
    
    def _init_tenant_db(self, db_path: str, plan: str):
        """Initialize the tenant's database schema"""
        conn = duckdb.connect(db_path)
        
        # Set resource limits per plan
        limits = {
            "free": {"memory": "256MB", "threads": 2},
            "pro":  {"memory": "1GB",   "threads": 4},
            "enterprise": {"memory": "4GB", "threads": 8},
        }
        limit = limits.get(plan, limits["free"])
        conn.execute(f"SET memory_limit = '{limit['memory']}'")
        conn.execute(f"SET threads = {limit['threads']}")
        
        # Create analytics tables
        conn.execute("""
            CREATE TABLE IF NOT EXISTS orders (
                order_id    BIGINT PRIMARY KEY,
                order_date  DATE NOT NULL,
                product     VARCHAR NOT NULL,
                category    VARCHAR NOT NULL,
                quantity    INTEGER NOT NULL,
                unit_price  DOUBLE NOT NULL,
                cost_price  DOUBLE NOT NULL,
                channel     VARCHAR NOT NULL,
                status      VARCHAR NOT NULL
            )
        """)
        
        # Pre-aggregated table (for faster common queries)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS daily_summary (
                report_date DATE PRIMARY KEY,
                revenue     DOUBLE,
                cost        DOUBLE,
                profit      DOUBLE,
                order_count INTEGER,
                avg_order   DOUBLE
            )
        """)
        
        conn.close()
    
    def get_connection(self, tenant_id: str) -> duckdb.DuckDBPyConnection:
        """Get a connection to the specified tenant's database"""
        result = self.meta_conn.execute(
            "SELECT db_path, status FROM tenants WHERE tenant_id = ?",
            [tenant_id]
        ).fetchone()
        
        if not result:
            raise ValueError(f"Tenant {tenant_id} not found")
        if result[1] != "active":
            raise ValueError(f"Tenant {tenant_id} is {result[1]}")
        
        return duckdb.connect(result[0])
    
    def cross_tenant_query(self, sql: str) -> list:
        """
        Cross-tenant query (admin only): 
        uses ATTACH to connect all active tenants
        """
        tenants = self.meta_conn.execute(
            "SELECT tenant_id, db_path FROM tenants WHERE status = 'active'"
        ).fetchall()
        
        # ATTACH all tenant databases
        admin_conn = duckdb.connect(":memory:")
        for tid, path in tenants:
            admin_conn.execute(f"ATTACH '{path}' AS {tid}")
        
        return admin_conn.execute(sql).fetchall()


# ══════════════════════════════════════════════════
# Usage Example
# ══════════════════════════════════════════════════
if __name__ == "__main__":
    manager = TenantDatabaseManager("/tmp/tenants_demo")
    
    # Create three tenants with different plans
    t1 = manager.create_tenant("Xiao Ming's Shop", "free")
    t2 = manager.create_tenant("Lao Wang Trading Co.", "pro")
    t3 = manager.create_tenant("Global Supply Chain Group", "enterprise")
    
    print(f"✅ Created 3 tenants:")
    print(f"   Free:     {t1}")
    print(f"   Pro:      {t2}")
    print(f"   Enterprise: {t3}")
    
    # Insert sample order data for tenant t2
    conn = manager.get_connection(t2)
    conn.execute("""
        INSERT INTO orders VALUES
            (1, '2026-05-01', 'Bluetooth Earbuds', 'Electronics', 120, 99.0, 40.0, 'Taobao', 'Completed'),
            (2, '2026-05-01', 'Power Bank',       'Electronics', 85,  79.0, 32.0, 'JD.com', 'Completed'),
            (3, '2026-05-02', 'Thermos',          'Home Goods',  200, 49.0, 20.0, 'Pinduoduo', 'Completed')
    """)
    conn.execute("""
        INSERT INTO daily_summary
        SELECT order_date, 
               SUM(quantity * unit_price) as revenue,
               SUM(quantity * cost_price) as cost,
               SUM(quantity * (unit_price - cost_price)) as profit,
               COUNT(DISTINCT order_id) as order_count,
               AVG(quantity * unit_price) as avg_order
        FROM orders GROUP BY order_date
    """)
    conn.close()
    
    # Query tenant t2's data
    conn = manager.get_connection(t2)
    result = conn.execute("""
        SELECT report_date, revenue, profit, 
               ROUND(profit/revenue*100, 1) as margin
        FROM daily_summary
    """).fetchdf()
    print(f"\n📊 Tenant {t2} Business Summary:")
    print(result)
    conn.close()
    
    # Cross-tenant admin query (using ATTACH)
    print("\n📈 All Tenant Summary:")
    admin_results = manager.cross_tenant_query("""
        SELECT 't2' as tenant_id, SUM(revenue) as total_revenue 
        FROM t2.daily_summary
        UNION ALL
        SELECT 't1', 0 FROM t1.daily_summary
    """)
    print(admin_results)

class HybridTenantManager:
    """
    Hybrid multi-tenant manager:
    - VIP tenants (Pro/Enterprise): independent database files
    - Regular tenants (Free): shared tables with tenant_id column
    """
    
    def __init__(self, data_dir: str = "/data/tenants"):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(parents=True, exist_ok=True)
        self.shared_db = str(self.data_dir / "_shared.duckdb")
        self._init_shared()
    
    def _init_shared(self):
        """Initialize the shared database (for regular tenants)"""
        conn = duckdb.connect(self.shared_db)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS shared_orders (
                tenant_id   VARCHAR NOT NULL,
                order_id    BIGINT NOT NULL,
                order_date  DATE NOT NULL,
                product     VARCHAR NOT NULL,
                quantity    INTEGER NOT NULL,
                amount      DOUBLE NOT NULL,
                PRIMARY KEY (tenant_id, order_id)
            )
        """)
        # Partitioned by tenant_id (DuckDB auto-optimizes)
        conn.execute("""
            CREATE TABLE IF NOT EXISTS shared_daily_summary (
                tenant_id   VARCHAR NOT NULL,
                report_date DATE NOT NULL,
                revenue     DOUBLE,
                order_count INTEGER,
                PRIMARY KEY (tenant_id, report_date)
            )
        """)
        conn.close()
    
    def query_with_isolation(self, tenant_id: str, sql: str) -> object:
        """
        Query with automatic tenant isolation.
        VIP tenants query their own DB, regular tenants 
        get automatic WHERE tenant_id= filters.
        """
        if self._is_vip_tenant(tenant_id):
            conn = duckdb.connect(str(self.data_dir / f"{tenant_id}.duckdb"))
        else:
            conn = duckdb.connect(self.shared_db)
            # Auto-inject tenant filter
            sql = f"SELECT * FROM ({sql}) sub WHERE sub.tenant_id = '{tenant_id}'"
        
        result = conn.execute(sql)
        conn.close()
        return result.fetchdf()
    
    def _is_vip_tenant(self, tenant_id: str) -> bool:
        """Determine if a tenant is VIP based on ID prefix"""
        return tenant_id.startswith("vip_")

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import duckdb
import time

app = FastAPI(title="DuckDB Multi-Tenant Analytics API")

# ─── Request/Response Models ───
class QueryRequest(BaseModel):
    tenant_id: str
    sql: str
    params: dict = {}

class QueryResponse(BaseModel):
    columns: list[str]
    rows: list[list]
    row_count: int
    execution_time_ms: float

class TenantInfo(BaseModel):
    tenant_id: str
    tenant_name: str
    plan: str

# ─── Dependency: Tenant Validation + DB Connection ───
def get_tenant_connection(tenant_id: str) -> duckdb.DuckDBPyConnection:
    """Validate tenant and return the corresponding database connection"""
    # In production, read from Redis or a metadata DB
    valid_tenants = {
        "t_demo_free": {"path": "/data/tenants/t_demo_free.duckdb", "plan": "free"},
        "t_demo_pro":  {"path": "/data/tenants/t_demo_pro.duckdb",  "plan": "pro"},
    }
    
    if tenant_id not in valid_tenants:
        raise HTTPException(status_code=404, detail="Tenant not found")
    
    info = valid_tenants[tenant_id]
    conn = duckdb.connect(info["path"])
    
    # Apply resource limits per plan
    if info["plan"] == "free":
        conn.execute("SET memory_limit = '256MB'")
        conn.execute("SET threads = 2")
    elif info["plan"] == "pro":
        conn.execute("SET memory_limit = '1GB'")
        conn.execute("SET threads = 4")
    
    return conn


# ─── API Endpoints ───

@app.post("/api/v1/query", response_model=QueryResponse)
async def run_query(req: QueryRequest):
    """Execute an SQL query with tenant isolation"""
    start = time.time()
    
    conn = get_tenant_connection(req.tenant_id)
    
    try:
        # Security: only allow SELECT queries
        sql_upper = req.sql.strip().upper()
        if not sql_upper.startswith("SELECT") and not sql_upper.startswith("WITH"):
            raise HTTPException(status_code=400, detail="Only SELECT queries allowed")
        
        # Block dangerous operations
        forbidden = ["DROP", "DELETE", "ALTER", "ATTACH", "DETACH",
                     "CREATE TABLE", "INSERT", "UPDATE"]
        for word in forbidden:
            if word in sql_upper:
                raise HTTPException(status_code=400, 
                    detail=f"Forbidden: {word} operations not allowed")
        
        result = conn.execute(req.sql, req.params)
        df = result.fetchdf()
        
        elapsed = (time.time() - start) * 1000
        
        return QueryResponse(
            columns=list(df.columns),
            rows=df.values.tolist(),
            row_count=len(df),
            execution_time_ms=round(elapsed, 2)
        )
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))
    finally:
        conn.close()


@app.get("/api/v1/tenant/{tenant_id}/info", response_model=TenantInfo)
async def get_tenant_info(tenant_id: str):
    """Get tenant information"""
    valid_tenants = {
        "t_demo_free": {"name": "Xiao Ming's Shop", "plan": "free"},
        "t_demo_pro":  {"name": "Lao Wang Trading Co.", "plan": "pro"},
    }
    if tenant_id not in valid_tenants:
        raise HTTPException(status_code=404, detail="Tenant not found")
    info = valid_tenants[tenant_id]
    return TenantInfo(
        tenant_id=tenant_id,
        tenant_name=info["name"],
        plan=info["plan"]
    )


@app.get("/api/v1/admin/total-revenue")
async def get_total_revenue():
    """
    Admin API: cross-tenant aggregation using ATTACH
    Note: Production use requires API Key authentication
    """
    admin_conn = duckdb.connect(":memory:")
    
    try:
        admin_conn.execute("ATTACH '/data/tenants/t_demo_free.duckdb' AS free_db")
        admin_conn.execute("ATTACH '/data/tenants/t_demo_pro.duckdb' AS pro_db")
        
        result = admin_conn.execute("""
            SELECT 'free' as tier, SUM(amount) as total_revenue 
            FROM free_db.orders
            UNION ALL
            SELECT 'pro', SUM(amount)
            FROM pro_db.orders
        """).fetchdf()
        
        return result.to_dict(orient="records")
    finally:
        admin_conn.close()


# ─── Startup ───
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)

-- Free plan: 256MB
SET memory_limit = '256MB';
SET threads = 2;

-- Pro plan: 1GB  
SET memory_limit = '1GB';
SET threads = 4;

-- Enterprise: 4GB
SET memory_limit = '4GB';
SET threads = 8;

#!/usr/bin/env python3
"""
Health check + auto-scaling script
Checks all tenant databases every 5 minutes
"""
import duckdb
from pathlib import Path
import json

def health_check(data_dir: str = "/data/tenants"):
    meta_path = Path(data_dir) / "_meta.duckdb"
    if not meta_path.exists():
        return {"status": "no_tenants"}
    
    conn = duckdb.connect(str(meta_path))
    
    result = conn.execute("""
        SELECT 
            tenant_id, tenant_name, plan, status,
            ROUND(data_size_mb, 1) as size_mb,
            CASE 
                WHEN data_size_mb > 500 THEN 'SCALE_UP'
                WHEN data_size_mb < 10 AND plan != 'free' THEN 'SCALE_DOWN'
                ELSE 'OK'
            END as action
        FROM tenants
        WHERE status = 'active'
    """).fetchdf()
    
    conn.close()
    return json.loads(result.to_json(orient="records"))

# Execute check
report = health_check()
print(f"🏥 Health check complete: {len(report)} active tenants")
for r in report:
    icon = "✅" if r['action'] == 'OK' else "⚠️"
    print(f"  {icon} {r['tenant_name']} ({r['plan']}) - {r['size_mb']}MB")