controller.py

import base64
from contextlib import asynccontextmanager
import logging
import time
import uuid
import os
import re
import asyncio  # Added for parallel batching
import multiprocessing  # Added for worker calculation
import hashlib # Added for Coalescing
import json # Added for Coalescing
import psutil # Added for Dynamic RAM calculation
from typing import List, Optional, Dict, Any, TypeVar, Generic  # Added Generic/TypeVar
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse

# --- FastAPI & Core ---
from fastapi import FastAPI, HTTPException, Depends
from fastapi.encoders import jsonable_encoder
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from fastapi.middleware.cors import CORSMiddleware
from starlette.concurrency import run_in_threadpool
from anyio import to_thread # Added for AnyIO 4.x concurrency tuning
from dotenv import load_dotenv
from pydantic import BaseModel

# --- Database Drivers & Pooling ---
from bson import ObjectId
import mysql.connector
from mysql.connector import pooling # For MySQL Pooling
import psycopg2
from psycopg2 import pool as pg_pool # For Postgres Pooling
from psycopg2.extras import RealDictCursor
from threading import Lock # To handle concurrency safely
import uvicorn

# --- Existing Services ---
from csv_analysis_service import execute_analysis_logic
from csv_chart_service import execute_python_code
from csv_metadata_service import CsvDataRequest, CsvInfoRequest, CsvInfoResponse, PythonExecutionRequest, PythonExecutionResponse, execute_python_logic, get_csv_basic_info, get_robust_csv_rows
from mongo_service import execute_mongo_operation, parse_query_input
from pydantic_csv_analysis_model import AnalysisRequest, AnalysisResponse
from pydantic_csv_charts_model import ChartExecutionPayload, ChartExecutionResponse
from pydantic_mongo_executor_model import ExecutorPayload, ExecutorResponse
from report_service import FileBoxProps, ReportRequest, execute_report_generation
from supabase_service import upload_bytes_to_supabase

# --- Configuration & Setup ---
load_dotenv()

logging.basicConfig(
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
    level=logging.INFO
)
logger = logging.getLogger("API_Controller")

def get_dynamic_thread_limit():
    """
    Calculates a safe thread limit based on available RAM and CPU cores.
    Returns an integer safe for to_thread.current_default_thread_limiter().total_tokens
    """
    try:
        # 1. Get System Resources
        total_ram_bytes = psutil.virtual_memory().total
        total_cores = multiprocessing.cpu_count()
        
        # 2. Identify how many Uvicorn Workers are running
        # Assuming production setup uses: max(1, total_cores - 2)
        num_workers = max(1, total_cores - 2)
        
        # 3. Calculate RAM available PER WORKER
        ram_per_worker = total_ram_bytes / num_workers
        
        # 4. Reserve Buffer (Keep 30% for OS/Python overhead, use 70% for threads)
        safe_ram_pool = ram_per_worker * 0.70
        
        # 5. Estimate Thread Cost (~8MB conservative safety margin)
        BYTES_PER_THREAD = 8 * 1024 * 1024 
        
        calculated_limit = int(safe_ram_pool / BYTES_PER_THREAD)
        
        # 6. Apply Reasonable Hard Caps (Min 100, Max 3000)
        final_limit = max(100, min(calculated_limit, 3000))
        
        logger.info(f"Dynamic Limit Config: {total_ram_bytes/(1024**3):.2f}GB RAM / {num_workers} Workers. Limit: {final_limit}")
        return final_limit

    except Exception as e:
        logger.warning(f"Failed to calculate dynamic threads ({e}). Fallback to 1000.")
        return 1000

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup logic: Calculate and set dynamic thread limit
    safe_limit = get_dynamic_thread_limit()
    to_thread.current_default_thread_limiter().total_tokens = safe_limit
    logger.info(f"Worker Process Started: Thread pool capacity set to {safe_limit}.")
    yield
    # Shutdown logic (if any) goes here

app = FastAPI(
    title="Unified Data Executor API (Mongo, SQL, CSV)", 
    lifespan=lifespan
)

# ==============================================================================
# HIGH-CONCURRENCY BATCH MODELS & STARTUP
# ==============================================================================
T = TypeVar("T")

class BatchRequest(BaseModel, Generic[T]):
    requests: List[T]

class BatchResponse(BaseModel, Generic[T]):
    responses: List[T]

# --- Directory Setup ---
CHART_DIR = "generated_charts"
os.makedirs(CHART_DIR, exist_ok=True)

# --- CORS ---
origins_env = os.getenv("ALLOWED_ORIGINS", "*")
ORIGINS = [origin.strip() for origin in origins_env.split(",")]

app.add_middleware(
    CORSMiddleware,
    allow_origins=ORIGINS,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --- Security ---
security = HTTPBearer()
API_SECRET_TOKEN = os.getenv("API_BEARER_TOKEN")

if not API_SECRET_TOKEN:
    logger.warning("WARNING: API_BEARER_TOKEN not set in .env file! Security is compromised.")

async def validate_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != API_SECRET_TOKEN:
        raise HTTPException(status_code=403, detail="Invalid Authentication Token")
    return credentials.credentials


# ==============================================================================
# PYDANTIC MODELS
# ==============================================================================

class SqlQueryRequest(BaseModel):
    database_url: str
    sql_query: str
    limit_rows: Optional[int] = 20
    limited: bool = False

class SqlQueryResponse(BaseModel):
    success: bool
    results: Optional[List[Dict[str, Any]]] = None
    columns: Optional[List[str]] = None
    rowCount: Optional[int] = 0
    executionTime: Optional[float] = 0.0
    error: Optional[str] = None
    request_id: str
    is_aggregate: bool = False
    limited: bool = False
    message: Optional[str] = None

class PgQueryRequest(BaseModel):
    database_url: str
    sql_query: str
    limit_rows: Optional[int] = 20
    limited: bool = False 

class PgQueryResponse(BaseModel):
    success: bool
    results: Optional[List[Dict[str, Any]]] = None
    columns: Optional[List[str]] = None
    rowCount: Optional[int] = 0
    executionTime: Optional[float] = 0.0
    error: Optional[str] = None
    request_id: str
    is_aggregate: bool = False
    limited: bool = False
    message: Optional[str] = None
    
# ==============================================================================
# CONNECTION POOL MANAGER
# ==============================================================================
class ConnectionPoolManager:
    def __init__(self):
        self._mysql_pools = {}
        self._pg_pools = {}
        self._lock = Lock() 

    def get_mysql_connection(self, db_url: str):
        with self._lock:
            if db_url not in self._mysql_pools:
                logger.info(f"Creating new MySQL pool for: {db_url}")
                parsed = urlparse(db_url)
                db_config = {
                    "user": parsed.username,
                    "password": parsed.password,
                    "host": parsed.hostname,
                    "port": parsed.port or 3306,
                    "database": parsed.path.lstrip("/"),
                    "connect_timeout": 5
                }
                self._mysql_pools[db_url] = pooling.MySQLConnectionPool(pool_name=str(uuid.uuid4()), pool_size=10, **db_config)
        return self._mysql_pools[db_url].get_connection()

    def get_postgres_connection(self, db_url: str):
        with self._lock:
            if db_url not in self._pg_pools:
                logger.info(f"Creating new Postgres pool for: {db_url}")
                parsed = urlparse(db_url)
                qs = parse_qs(parsed.query)
                sslmode = qs.get('sslmode', ['require'])[0] if 'sslmode' in qs else 'prefer'
                db_config = {
                    "host": parsed.hostname,
                    "port": parsed.port or 5432,
                    "database": parsed.path.lstrip("/"),
                    "user": parsed.username,
                    "password": parsed.password,
                    "sslmode": sslmode,
                    "connect_timeout": 5
                }
                self._pg_pools[db_url] = pg_pool.ThreadedConnectionPool(1, 10, **db_config)
        return self._pg_pools[db_url].getconn()

    def return_postgres_connection(self, db_url, conn, close=False):
        """
        Returns connection to pool. 
        If close=True, the connection is discarded (used for dead connections).
        """
        if db_url in self._pg_pools and conn:
            self._pg_pools[db_url].putconn(conn, close=close)

pool_manager = ConnectionPoolManager()


# ==============================================================================
# REQUEST COALESCER (SINGLEFLIGHT)
# ==============================================================================
class RequestCoalescer:
    def __init__(self):
        self._active_requests: Dict[str, asyncio.Future] = {}
        self._lock = asyncio.Lock()

    def _generate_key(self, prefix: str, data: dict) -> str:
        json_str = json.dumps(data, sort_keys=True, default=str)
        raw_str = f"{prefix}:{json_str}"
        return hashlib.md5(raw_str.encode()).hexdigest()

    async def execute(self, prefix: str, unique_params: dict, func, *args, **kwargs):
        key = self._generate_key(prefix, unique_params)

        async with self._lock:
            if key in self._active_requests:
                return await self._active_requests[key]
            
            loop = asyncio.get_running_loop()
            future = loop.create_future()
            self._active_requests[key] = future

        try:
            # Execute the function (run_in_threadpool) with args/kwargs
            result = await func(*args, **kwargs)
            
            if not future.done():
                future.set_result(result)
            return result
            
        except Exception as e:
            if not future.done():
                future.set_exception(e)
            raise e
        finally:
            async with self._lock:
                if key in self._active_requests:
                    del self._active_requests[key]

coalescer = RequestCoalescer()


# ==============================================================================
# DB EXECUTION LOGIC
# ==============================================================================

def is_aggregate_query(query: str) -> bool:
    query_lower = query.lower()
    aggregate_patterns = [
        r'\bcount\s*\(', r'\bsum\s*\(', r'\bavg\s*\(', 
        r'\bmin\s*\(', r'\bmax\s*\(', r'\bgroup\s+by\b',
        r'\bdistinct\b', r'\bhaving\b'
    ]
    for pattern in aggregate_patterns:
        if re.search(pattern, query_lower):
            return True
    return False

def normalize_mysql_uri(uri: str) -> str:
    try:
        parsed_uri = urlparse(uri)
        query_params = parse_qs(parsed_uri.query)
        query_params.pop('ssl-mode', None)
        new_query = urlencode(query_params, doseq=True)
        parsed_uri = parsed_uri._replace(query=new_query)
        return urlunparse(parsed_uri)
    except Exception:
        return uri

def normalize_postgres_uri(uri: str) -> str:
    try:
        parsed_uri = urlparse(uri)
        if parsed_uri.scheme == 'postgres':
            parsed_uri = parsed_uri._replace(scheme='postgresql')
        return urlunparse(parsed_uri)
    except Exception:
        return uri

def _run_mysql_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
    start_time = time.time()
    connection = None
    cursor = None
    response = {"success": False, "results": None, "columns": None, "rowCount": 0, "executionTime": 0.0, "error": None, "is_aggregate": False, "limited": False, "message": ""}
    
    # Flag for bad connections
    connection_broken = False

    try:
        connection = pool_manager.get_mysql_connection(db_url)
        cursor = connection.cursor(dictionary=True)
        
        clean_query = sql_query.strip()
        query_lower = clean_query.lower()

        if not query_lower.startswith("select"):
            cursor.execute(clean_query)
            connection.commit()
            response.update({"success": True, "message": "Query executed successfully (Non-SELECT)."})
            return response

        if not limited:
            cursor.execute(clean_query)
            results = cursor.fetchall()
            response["message"] = f"Raw query executed. Returned {len(results)} row(s)."
            response["limited"] = False
            response["is_aggregate"] = is_aggregate_query(clean_query)
        else:
            if is_aggregate_query(clean_query):
                cursor.execute(clean_query)
                results = cursor.fetchall()
                response["message"] = f"Aggregate query completed. Returned {len(results)} row(s)."
                response["is_aggregate"] = True
            else:
                final_query = clean_query.rstrip(';').strip()
                if not re.search(r'\blimit\s+\d+', query_lower):
                    final_query = f"{final_query} LIMIT {max_rows}"
                
                cursor.execute(final_query)
                results = cursor.fetchall()
                is_limited_result = (len(results) == max_rows)
                response["message"] = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
                response["limited"] = is_limited_result
                response["is_aggregate"] = False

        columns = [col[0] for col in cursor.description] if cursor.description else []
        response.update({"success": True, "results": jsonable_encoder(results), "columns": columns, "rowCount": len(results), "executionTime": time.time() - start_time})
        return response

    except Exception as e:
        # Detect broken pipe or connection lost errors in MySQL
        err_str = str(e).lower()
        if "lost connection" in err_str or "gone away" in err_str or isinstance(e, mysql.connector.errors.OperationalError):
            connection_broken = True
        
        response["error"] = str(e)
        return response
    finally:
        try:
            if cursor: cursor.close()
        except: pass
        
        if connection:
            if connection_broken:
                # Discard dead connection (do NOT return to pool)
                try: connection.close() 
                except: pass
            else:
                # Return healthy connection to pool
                connection.close() 

def _run_postgres_synchronously(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
    start_time = time.time()
    connection = None
    cursor = None
    response = {"success": False, "results": None, "columns": None, "rowCount": 0, "executionTime": 0.0, "error": None, "is_aggregate": False, "limited": False, "message": ""}
    
    # Flag to determine if connection is dead
    connection_broken = False 

    try:
        connection = pool_manager.get_postgres_connection(db_url)
        cursor = connection.cursor(cursor_factory=RealDictCursor)
        clean_query = sql_query.strip()
        query_lower = clean_query.lower()
        
        if not query_lower.startswith(("select", "show", "explain", "with")):
            cursor.execute(clean_query)
            connection.commit()
            response.update({"success": True, "message": "Query executed successfully (Non-SELECT)."})
            return response
            
        if not limited:
            cursor.execute(clean_query)
            results = cursor.fetchall()
            response["message"] = f"Raw query executed. Returned {len(results)} row(s)."
            response["limited"] = False
            response["is_aggregate"] = is_aggregate_query(clean_query)
        else:
            if is_aggregate_query(clean_query):
                cursor.execute(clean_query)
                results = cursor.fetchall()
                response["message"] = f"Aggregate query completed. Returned {len(results)} row(s)."
                response["is_aggregate"] = True
                response["limited"] = False
            else:
                final_query = clean_query.rstrip(';').strip()
                if not re.search(r'\blimit\s+\d+', query_lower):
                    final_query = f"{final_query} LIMIT {max_rows}"
                cursor.execute(final_query)
                results = cursor.fetchall()
                is_limited_result = (len(results) == max_rows)
                response["message"] = f"Showing first {max_rows} rows only." if is_limited_result else f"Returned {len(results)} rows."
                response["limited"] = is_limited_result
                response["is_aggregate"] = False

        columns = [desc[0] for desc in cursor.description] if cursor.description else []
        clean_results = jsonable_encoder(results, custom_encoder={uuid.UUID: str, ObjectId: str})
        response.update({"success": True, "results": clean_results, "columns": columns, "rowCount": len(results), "executionTime": time.time() - start_time})
        return response
    except Exception as e:
        # Detect fatal connection errors
        err_msg = str(e).lower()
        if "closed" in err_msg or "terminat" in err_msg or isinstance(e, (psycopg2.InterfaceError, psycopg2.OperationalError)):
            connection_broken = True 
        
        if connection and not connection_broken:
            try:
                connection.rollback()
            except Exception:
                connection_broken = True 

        response["error"] = str(e)
        return response
    finally:
        try: 
            if cursor: cursor.close()
        except: pass
        
        if connection: 
            # If broken, set close=True to discard it from pool
            pool_manager.return_postgres_connection(db_url, connection, close=connection_broken)


# ==============================================================================
# API ROUTES
# ==============================================================================

@app.post("/api/execute_mongo", response_model=ExecutorResponse)
async def execute_mongo_endpoint(payload: ExecutorPayload, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    start_time = time.time()
    try:
        # Parse logic is fast, can stay outside threadpool
        parsed_query = parse_query_input(payload.generated_query)
        
        # --- COALESCING MAGIC ---
        # Create a unique signature for this request
        unique_params = {
            "uri": payload.mongo_uri,
            "db": payload.db_name,
            "col": payload.collection_name,
            "q": parsed_query, # parsed_query is a Dict or List, json.dumps handles it
            "lim": payload.limited,
            "lrows": payload.limit_rows
        }

        # Use the coalescer to prevent duplicate simultaneous DB hits
        result_data = await coalescer.execute(
            "mongo",                 # Prefix
            unique_params,           # Unique params dict
            run_in_threadpool,       # Runner
            execute_mongo_operation, # The function to run
            # Arguments for execute_mongo_operation:
            mongo_uri=payload.mongo_uri,
            db_name=payload.db_name,
            collection_name=payload.collection_name,
            query=parsed_query,
            limited=payload.limited,
            limit_rows=payload.limit_rows
        )
        # ------------------------

        return ExecutorResponse(
            status="success", 
            count=len(result_data), 
            data=jsonable_encoder(result_data, custom_encoder={ObjectId: str}), 
            duration_seconds=round(time.time() - start_time, 4), 
            request_id=request_id
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/execute_chart", response_model=ChartExecutionResponse)
async def execute_chart_endpoint(payload: ChartExecutionPayload, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        image_bytes, error_msg, logs = await run_in_threadpool(execute_python_code, code=payload.code, csv_url=payload.csv_url)
        
        if error_msg: 
            return ChartExecutionResponse(status="error", error=error_msg, output_log=logs, request_id=request_id)

        if payload.return_base64:
            base64_str = base64.b64encode(image_bytes).decode('utf-8')
            return ChartExecutionResponse(status="success", base64_image=base64_str, output_log=logs, request_id=request_id)
        else:
            unique_name = f"{uuid.uuid4()}.png"
            public_url = await run_in_threadpool(upload_bytes_to_supabase, image_bytes=image_bytes, file_name=unique_name, chat_id=payload.chat_id)
            return ChartExecutionResponse(status="success", image_url=public_url, output_log=logs, request_id=request_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.post("/api/execute_sql_query", response_model=SqlQueryResponse)
async def execute_mysql_endpoint(query: SqlQueryRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        normalized_url = normalize_mysql_uri(query.database_url)
        limit_val = query.limit_rows if query.limit_rows is not None else 20
        
        # Unique ID for Coalescing
        unique_params = {
            "db": normalized_url,
            "q": query.sql_query,
            "l": limit_val,
            "lim": query.limited
        }

        # FIXED: Pass _run_mysql_synchronously as the first POSITIONAL argument after run_in_threadpool
        result_dict = await coalescer.execute(
            "mysql",                  # Prefix
            unique_params,            # Unique Params
            run_in_threadpool,        # Runner
            _run_mysql_synchronously, # Arg 1 (The Function)
            db_url=normalized_url,    # Kwargs
            sql_query=query.sql_query,
            max_rows=limit_val,
            limited=query.limited
        )

        final_result = result_dict.copy()
        final_result["request_id"] = request_id
        return SqlQueryResponse(**final_result)

    except Exception as e:
        # If run_in_threadpool fails (e.g. TypeError) it ends up here
        logger.error(f"MySQL Endpoint Error: {str(e)}")
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})

@app.post("/api/execute_postgres_query", response_model=PgQueryResponse)
async def execute_postgres_endpoint(query: PgQueryRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        clean_url = normalize_postgres_uri(query.database_url)
        limit_val = query.limit_rows if query.limit_rows is not None else 20
        
        unique_params = {
            "db": clean_url,
            "q": query.sql_query,
            "l": limit_val,
            "lim": query.limited
        }

        # FIXED: Pass _run_postgres_synchronously as the first POSITIONAL argument
        result_dict = await coalescer.execute(
            "postgres",
            unique_params,
            run_in_threadpool,
            _run_postgres_synchronously, # Arg 1 (The Function)
            db_url=clean_url,
            sql_query=query.sql_query,
            max_rows=limit_val,
            limited=query.limited
        )

        final_result = result_dict.copy()
        final_result["request_id"] = request_id
        return PgQueryResponse(**final_result)
    except Exception as e:
        logger.error(f"Postgres Endpoint Error: {str(e)}")
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})

# ... (CSV and Report endpoints remain unchanged) ...
@app.post("/api/execute_csv_analysis", response_model=AnalysisResponse)
async def execute_analysis_endpoint(payload: AnalysisRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        result = await run_in_threadpool(execute_analysis_logic, code=payload.code, csv_url=payload.csv_url)
        return AnalysisResponse(success=result["success"], output_log=result["output_log"], results=result["results"], error=result["error"], request_id=request_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
@app.post("/api/generate_report", response_model=FileBoxProps)
async def generate_report_endpoint(payload: ReportRequest, token: str = Depends(validate_token)):
    try:
        result = await execute_report_generation(code=payload.code, csv_url=payload.csv_url, chat_id=payload.chat_id)
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
@app.post("/api/get_csv_info", response_model=CsvInfoResponse)
async def get_csv_info_endpoint(payload: CsvInfoRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    start_time = time.time()
    try:
        info_result = await run_in_threadpool(get_csv_basic_info, csv_path=payload.csv_url)
        if "error" in info_result:
            return CsvInfoResponse(success=False, error=info_result["error"], request_id=request_id, duration=time.time() - start_time)
        return CsvInfoResponse(success=True, data=info_result, request_id=request_id, duration=time.time() - start_time)
    except Exception as e:
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})
        
@app.post("/api/csv_data")
async def get_csv_data_endpoint(payload: CsvDataRequest, token: str = Depends(validate_token)):
    try:
        result = await run_in_threadpool(get_robust_csv_rows, csv_url=payload.csv_url)
        if isinstance(result, dict) and "error" in result: raise HTTPException(status_code=400, detail=result["error"])
        return result
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")

@app.post("/api/execute_python", response_model=PythonExecutionResponse)
async def execute_python_endpoint(payload: PythonExecutionRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        execution_result = await run_in_threadpool(execute_python_logic, code=payload.code, custom_context=payload.context)
        return PythonExecutionResponse(success=execution_result['error'] is None, output=execution_result['output'], result=jsonable_encoder(execution_result['result']), isStructured=execution_result['isStructured'], error=execution_result['error'], request_id=request_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})

# --- Batch Handlers ---
async def batch_parallel_handler(func, requests: List[Any], token: str):
    tasks = [func(req, token) for req in requests]
    results = await asyncio.gather(*tasks, return_exceptions=True)
    return [res if not isinstance(res, Exception) else {"success": False, "error": str(res)} for res in results]

@app.post("/api/batch/execute_sql_query", response_model=BatchResponse[SqlQueryResponse])
async def batch_execute_sql(payload: BatchRequest[SqlQueryRequest], token: str = Depends(validate_token)):
    responses = await batch_parallel_handler(execute_mysql_endpoint, payload.requests, token)
    return BatchResponse(responses=responses)

@app.post("/api/batch/execute_postgres_query", response_model=BatchResponse[PgQueryResponse])
async def batch_execute_pg(payload: BatchRequest[PgQueryRequest], token: str = Depends(validate_token)):
    responses = await batch_parallel_handler(execute_postgres_endpoint, payload.requests, token)
    return BatchResponse(responses=responses)

@app.post("/api/batch/execute_mongo", response_model=BatchResponse[ExecutorResponse])
async def batch_execute_mongo(payload: BatchRequest[ExecutorPayload], token: str = Depends(validate_token)):
    responses = await batch_parallel_handler(execute_mongo_endpoint, payload.requests, token)
    return BatchResponse(responses=responses)

# --- Test Endpoint ---

# ==============================================================================
# KEEP-ALIVE ROUTE
# ==============================================================================
@app.get("/")
async def root():
    return {"message": "Python Code Execution Server is running"}

@app.get("/ping")
async def ping():
    return {"message": "I am alive!"}


# ==============================================================================
# HIGH PERFORMANCE SERVER EXECUTION
# ==============================================================================

if __name__ == "__main__":
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 7860))
    
    # Calculate workers (save 2 cores for system overhead)
    # Ensure at least 1 worker exists
    num_workers = max(1, multiprocessing.cpu_count() - 2)
    
    print(f"Starting production server on {host}:{port} with {num_workers} workers...")
    print("Using 'asyncio' loop to prevent Pandas/Numpy segfaults.")

    uvicorn.run(
        "controller:app",  
        host=host, 
        port=port, 
        workers=num_workers,
        loop="asyncio",
    )