"""
DuckDB DataHub — Unified Data Access Layer
==========================================
Manages unified querying across PostgreSQL, MySQL, CSV, Parquet, and HTTP APIs.

Installation:
 pip install duckdb pandas pyarrow

Usage:
 hub = DataHub()
 hub.attach_pg("orders_db", "host=localhost dbname=orders user=analyst")
 hub.attach_csv("warehouse", "/data/warehouse/")
 df = hub.query("SELECT * FROM orders_db.orders LIMIT 10")
 hub.close()
"""

import duckdb
import pandas as pd
from typing import Optional, Dict, List, Any
from pathlib import Path


class DataHub:
 """Production-grade DuckDB data access center"""
 
 def __init__(self, memory_limit: str = "50%", threads: int = 0):
 """
 Args:
 memory_limit: DuckDB memory limit, e.g., "50%", "8GB", "4g"
 threads: Parallel thread count, 0 = auto-detect CPU cores
 """
 self.con = duckdb.connect(
 ":memory:",
 config={
 "memory_limit": memory_limit,
 "threads": threads,
 "enable_object_cache": True,
 "default_order": "DESC",
 }
 )
 self._initialized = False
 self._attached_sources: Dict[str, dict] = {}
 
 def init_extensions(self):
 """Initialize all extensions (idempotent operation)"""
 if self._initialized:
 return
 
 extensions = [
 "postgres_scanner",
 "mysql_scanner",
 "httpfs",
 "json",
 "sqlite",
 "parquet",
 ]
 
 for ext in extensions:
 try:
 self.con.execute(f"INSTALL {ext}; LOAD {ext};")
 except Exception:
 pass # Extension may already be installed or unavailable
 
 self._initialized = True
 
 # ========== Data Source Attachment ==========
 
 def attach_pg(self, name: str, conn_str: str, schema: str = "public") -> None:
 """Attach a PostgreSQL database"""
 self.init_extensions()
 self.con.execute(f"""
 ATTACH '{conn_str}' AS {name} (TYPE POSTGRES, SCHEMA '{schema}');
 """)
 self._attached_sources[name] = {
 "type": "postgres",
 "conn_str": conn_str,
 "schema": schema,
 }
 print(f"✅ Attached PostgreSQL: {name}")
 
 def attach_mysql(self, name: str, conn_str: str, schema: str = "public") -> None:
 """Attach a MySQL database"""
 self.init_extensions()
 self.con.execute(f"""
 ATTACH '{conn_str}' AS {name} (TYPE MYSQL, SCHEMA '{schema}');
 """)
 self._attached_sources[name] = {
 "type": "mysql",
 "conn_str": conn_str,
 "schema": schema,
 }
 print(f"✅ Attached MySQL: {name}")
 
 def attach_csv_dir(self, name: str, dir_path: str, 
 pattern: str = "*.csv", header: bool = True) -> None:
 """Attach a CSV directory (auto-scans all matching files)"""
 self.con.execute(f"""
 ATTACH '{dir_path}' AS {name} (TYPE CSV, HEADER {header}, 
 READ_ALL_COLUMNS TRUE, 
 FILE_PATTERN '{pattern}');
 """)
 self._attached_sources[name] = {
 "type": "csv_dir",
 "path": dir_path,
 "pattern": pattern,
 }
 print(f"✅ Attached CSV directory: {name} ({dir_path})")
 
 def attach_parquet_dir(self, name: str, dir_path: str) -> None:
 """Attach a Parquet directory"""
 self.con.execute(f"""
 ATTACH '{dir_path}' AS {name} (TYPE PARQUET);
 """)
 self._attached_sources[name] = {
 "type": "parquet_dir",
 "path": dir_path,
 }
 print(f"✅ Attached Parquet directory: {name} ({dir_path})")
 
 def attach_http_json(self, name: str, url: str) -> None:
 """Attach an HTTP JSON API"""
 self.init_extensions()
 self.con.execute(f"""
 CREATE OR REPLACE VIEW {name} AS 
 SELECT * FROM read_json_auto('{url}');
 """)
 self._attached_sources[name] = {
 "type": "http_json",
 "url": url,
 }
 print(f"✅ Attached HTTP API: {name} ({url})")
 
 def attach_http_csv(self, name: str, url: str) -> None:
 """Attach an HTTP CSV API"""
 self.init_extensions()
 self.con.execute(f"""
 CREATE OR REPLACE VIEW {name} AS 
 SELECT * FROM read_csv_auto('{url}');
 """)
 self._attached_sources[name] = {
 "type": "http_csv",
 "url": url,
 }
 print(f"✅ Attached HTTP CSV: {name} ({url})")
 
 def attach_local_csv(self, name: str, file_path: str) -> None:
 """Attach a single CSV file"""
 self.con.execute(f"""
 CREATE OR REPLACE VIEW {name} AS 
 SELECT * FROM read_csv_auto('{file_path}');
 """)
 self._attached_sources[name] = {
 "type": "local_csv",
 "path": file_path,
 }
 print(f"✅ Attached CSV file: {name} ({file_path})")
 
 # ========== Query Interface ==========
 
 def query(self, sql: str, params: tuple = ()) -> pd.DataFrame:
 """Execute SQL query and return DataFrame"""
 result = self.con.execute(sql, params)
 if result.description:
 df = result.fetchdf()
 return df
 return pd.DataFrame()
 
 def query_raw(self, sql: str, params: tuple = ()) -> duckdb.DuckDBPyRelation:
 """Execute SQL query and return raw Relation (supports lazy reading)"""
 return self.con.execute(sql, params)
 
 def describe(self, table_ref: str) -> pd.DataFrame:
 """View schema information for a table"""
 return self.con.execute(f"DESCRIBE {table_ref}").fetchdf()
 
 def table_names(self, schema: str = "main") -> List[str]:
 """List all tables in a given schema"""
 result = self.con.execute(f"""
 SELECT table_name FROM information_schema.tables 
 WHERE table_schema = '{schema}'
 """).fetchdf()
 return result["table_name"].tolist()
 
 # ========== Cache & Optimization ==========
 
 def create_cached_table(self, name: str, sql: str) -> None:
 """Cache query results as a local table (speeds up repeated queries)"""
 self.con.execute(f"CREATE OR REPLACE TABLE {name} AS {sql}")
 print(f"✅ Created cached table: {name}")
 
 def drop_source(self, name: str) -> None:
 """Detach a data source"""
 self.con.execute(f"DETACH {name}")
 self._attached_sources.pop(name, None)
 print(f"🗑️ Detached: {name}")
 
 # ========== Diagnostic Tools ==========
 
 def explain(self, sql: str) -> pd.DataFrame:
 """Get SQL execution plan"""
 return self.con.execute(f"EXPLAIN {sql}").fetchdf()
 
 def explain_analyze(self, sql: str) -> pd.DataFrame:
 """Get actual execution plan with timing"""
 return self.con.execute(f"EXPLAIN ANALYZE {sql}").fetchdf()
 
 def show_attached(self) -> pd.DataFrame:
 """Show all attached data sources"""
 return self.con.execute("SHOW DATABASES").fetchdf()
 
 # ========== Lifecycle ==========
 
 def close(self) -> None:
 """Close connections and release resources"""
 self.con.close()
 print("🔒 DataHub closed")
 
 def __enter__(self):
 return self
 
 def __exit__(self, exc_type, exc_val, exc_tb):
 self.close()

# Use context manager for automatic connection cleanup
with DataHub(memory_limit="4GB") as hub:
 # Attach data sources
 hub.attach_pg("pg_orders", "host=192.168.1.10 dbname=ecommerce user=analyst password=***")
 hub.attach_csv_dir("csv_reports", "/data/reports/")
 hub.attach_http_json("exchange_rates", "https://api.exchangerate.host/latest?base=CNY")
 
 # Cross-source query
 df = hub.query("""
 WITH daily_sales AS (
 SELECT 
 o.order_date,
 c.category_name,
 SUM(o.amount) AS total_amount
 FROM pg_orders.orders o
 JOIN pg_orders.categories c ON o.category_id = c.category_id
 WHERE o.order_date >= '2026-01-01'
 GROUP BY o.order_date, c.category_name
 )
 SELECT 
 order_date,
 category_name,
 total_amount,
 total_amount * rates.usd_rate AS total_usd
 FROM daily_sales
 CROSS JOIN exchange_rates rates
 ORDER BY total_usd DESC
 LIMIT 20
 """)
 
 print(df.head())
 
 # Check execution plan to confirm predicate pushdown
 plan = hub.explain_analyze("""
 SELECT category_name, SUM(amount) 
 FROM pg_orders.orders 
 WHERE order_date = '2026-06-01' 
 GROUP BY category_name
 """)
 print(plan)

-- ✅ GOOD: Pushes down to PostgreSQL
SELECT category, SUM(amount) 
FROM pg_db.orders 
WHERE order_date = '2026-06-01' 
GROUP BY category;

-- ❌ BAD: Pulls entire table then filters
SELECT * FROM pg_db.orders;

EXPLAIN ANALYZE 
SELECT category, SUM(amount) 
FROM pg_db.orders 
WHERE order_date = '2026-06-01';

-- First time: create cache
CREATE TABLE daily_category_sales AS
SELECT 
 DATE_TRUNC('day', order_date) AS sale_date,
 category_name,
 SUM(amount) AS total_amount,
 COUNT(*) AS order_count
FROM pg_db.orders
WHERE order_date >= '2026-01-01'
GROUP BY 1, 2;

-- Later: query local table (millisecond response)
SELECT * FROM daily_category_sales 
WHERE sale_date = '2026-06-24';

class DataHubPool:
 """Data source connection pool"""
 
 def __init__(self, max_connections: int = 5):
 self.hubs = []
 self.max = max_connections
 
 def get_hub(self) -> DataHub:
 if self.hubs:
 return self.hubs.pop()
 return DataHub()
 
 def release(self, hub: DataHub):
 if len(self.hubs) < self.max:
 self.hubs.append(hub)
 else:
 hub.close()

hub = DataHub(memory_limit="8GB", threads=8)

# Key configuration options
config = {
 "memory_limit": "8GB", # Maximum memory
 "threads": 8, # Parallelism
 "max_threads": 16, # Maximum thread count
 "temp_directory": "/tmp/duckdb", # Temp directory (for large queries)
 "enable_object_cache": True, # Enable object cache
 "allow_unsigned_extensions": True,
}

# 1. Initialize Hub
hub = DataHub(memory_limit="16GB")

# 2. Attach all data sources
hub.attach_pg("orders", "host=db-host dbname=ecommerce user=analyst")
hub.attach_pg("crm", "host=crm-host dbname=customer_db user=analyst")
hub.attach_csv_dir("marketing", "/data/marketing/")
hub.attach_parquet_dir("warehouse", "/data/warehouse/parquet/")
hub.attach_http_json("competitors", "https://api.competitor.com/prices")

# 3. Create caching layer (daily refresh)
hub.create_cached_table("monthly_summary", """
 SELECT 
 DATE_TRUNC('month', o.order_date) AS month,
 c.category_name,
 COUNT(*) AS orders,
 SUM(o.amount) AS revenue,
 AVG(o.amount) AS avg_order_value
 FROM orders.orders o
 JOIN orders.categories c ON o.category_id = c.category_id
 GROUP BY 1, 2
""")

# 4. Business users only need to run SQL
df = hub.query("""
 SELECT month, category_name, revenue, 
 revenue / LAG(revenue) OVER (ORDER BY month) - 1 AS mom_growth
 FROM monthly_summary
 WHERE category_name = 'Electronics'
 ORDER BY month DESC
""")