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")