controller.py

import base64
from contextlib import asynccontextmanager
import logging
import time
import uuid
import os
import re
import asyncio
import multiprocessing
import hashlib
import json
import psutil
from typing import List, Optional, Dict, Any, TypeVar, Generic
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse
from threading import Lock

# --- 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
from dotenv import load_dotenv
from pydantic import BaseModel

# --- Async Database Drivers ---
import asyncpg
import asyncmy
from asyncmy.cursors import DictCursor

# --- Database Drivers (Kept for compatibility/types if needed, but unused in SQL routes) ---
from bson import ObjectId
import mysql.connector
import psycopg2
from psycopg2.extras import RealDictCursor
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")

# ==============================================================================
# ASYNC CONNECTION POOL MANAGER
# ==============================================================================
class AsyncPoolManager:
    def __init__(self):
        self._pg_pools: Dict[str, asyncpg.Pool] = {}
        self._mysql_pools: Dict[str, asyncmy.Pool] = {}
        self._lock = asyncio.Lock()

    async def get_pg_pool(self, db_url: str) -> asyncpg.Pool:
        async with self._lock:
            if db_url not in self._pg_pools:
                logger.info(f"Creating new AsyncPG pool for: {db_url[:20]}...")
                try:
                    # asyncpg handles Keepalives and SSL automatically better than psycopg2
                    pool = await asyncpg.create_pool(
                        dsn=db_url,
                        min_size=1,
                        max_size=20,
                        max_inactive_connection_lifetime=300,
                        command_timeout=60
                    )
                    self._pg_pools[db_url] = pool
                except Exception as e:
                    logger.error(f"Failed to create PG pool: {e}")
                    raise e
        return self._pg_pools[db_url]

    async def get_mysql_pool(self, db_url: str) -> asyncmy.Pool:
        async with self._lock:
            if db_url not in self._mysql_pools:
                logger.info(f"Creating new AsyncMy pool for: {db_url[:20]}...")
                parsed = urlparse(db_url)
                # Check for SSL requirement in query params
                qs = parse_qs(parsed.query)
                # AsyncMy often auto-negotiates SSL, but we can enforce it if needed based on qs
                
                try:
                    pool = await asyncmy.create_pool(
                        user=parsed.username,
                        password=parsed.password,
                        host=parsed.hostname,
                        port=parsed.port or 3306,
                        db=parsed.path.lstrip("/"),
                        minsize=1,
                        maxsize=20,
                        autocommit=True,
                        pool_recycle=280, # Recycle before default 8hr timeout to prevent "Gone Away"
                    )
                    self._mysql_pools[db_url] = pool
                except Exception as e:
                    logger.error(f"Failed to create MySQL pool: {e}")
                    raise e
        return self._mysql_pools[db_url]

    async def close_all(self):
        logger.info("Closing all async database pools...")
        for url, pool in self._pg_pools.items():
            await pool.close()
        for url, pool in self._mysql_pools.items():
            pool.close()
            await pool.wait_closed()

# Initialize Global Async Manager
async_pool_manager = AsyncPoolManager()

def get_dynamic_thread_limit():
    """Calculates a safe thread limit based on available RAM."""
    try:
        total_ram_bytes = psutil.virtual_memory().total
        total_cores = multiprocessing.cpu_count()
        num_workers = max(1, total_cores - 2)
        ram_per_worker = total_ram_bytes / num_workers
        safe_ram_pool = ram_per_worker * 0.70
        BYTES_PER_THREAD = 8 * 1024 * 1024 
        calculated_limit = int(safe_ram_pool / BYTES_PER_THREAD)
        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
    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}. Async Drivers Ready.")
    yield
    # Shutdown
    await async_pool_manager.close_all()

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

# ==============================================================================
# MODELS & MIDDLEWARE
# ==============================================================================
T = TypeVar("T")
class BatchRequest(BaseModel, Generic[T]): requests: List[T]
class BatchResponse(BaseModel, Generic[T]): responses: List[T]

CHART_DIR = "generated_charts"
os.makedirs(CHART_DIR, exist_ok=True)

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 = HTTPBearer()
API_SECRET_TOKEN = os.getenv("API_BEARER_TOKEN")

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

# ==============================================================================
# REQUEST COALESCER
# ==============================================================================
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]
            
            future = asyncio.get_running_loop().create_future()
            self._active_requests[key] = future

        try:
            # Works for both async functions and run_in_threadpool
            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()

# ==============================================================================
# ASYNC DB 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

# --- ASYNC MYSQL EXECUTOR ---
async def _execute_async_mysql(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
    start_time = time.time()
    try:
        # Get pool
        pool = await async_pool_manager.get_mysql_pool(db_url)
        
        async with pool.acquire() as conn:
            async with conn.cursor(cursor=DictCursor) as cursor:
                clean_query = sql_query.strip()
                query_lower = clean_query.lower()

                if not query_lower.startswith("select"):
                    await cursor.execute(clean_query)
                    return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}

                final_query = clean_query
                is_aggregate = is_aggregate_query(clean_query)
                is_limited_result = False
                message = ""

                if not limited:
                    message = f"Raw query executed."
                else:
                    if is_aggregate:
                        message = f"Aggregate query completed."
                    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}"
                        is_limited_result = True
                        message = f"Showing first {max_rows} rows."

                await cursor.execute(final_query)
                results = await cursor.fetchall()
                
                # Check actual length
                if is_limited_result and len(results) < max_rows:
                    is_limited_result = False
                    message = f"Returned {len(results)} rows."
                elif is_limited_result:
                    message = f"Showing first {len(results)} rows."

                columns = [col[0] for col in cursor.description] if cursor.description else []
                
                return {
                    "success": True, "results": jsonable_encoder(results), "columns": columns, 
                    "rowCount": len(results), "executionTime": time.time() - start_time,
                    "is_aggregate": is_aggregate, "limited": is_limited_result, "message": message, "error": None
                }
    except Exception as e:
        return {"success": False, "error": str(e), "executionTime": 0.0}

# --- ASYNC POSTGRES EXECUTOR ---
async def _execute_async_postgres(db_url: str, sql_query: str, max_rows: int = 20, limited: bool = False) -> dict:
    start_time = time.time()
    try:
        pool = await async_pool_manager.get_pg_pool(db_url)
        
        async with pool.acquire() as conn:
            clean_query = sql_query.strip()
            query_lower = clean_query.lower()

            if not query_lower.startswith(("select", "show", "explain", "with")):
                await conn.execute(clean_query)
                return {"success": True, "message": "Query executed successfully (Non-SELECT).", "executionTime": time.time() - start_time}

            final_query = clean_query
            is_aggregate = is_aggregate_query(clean_query)
            is_limited_result = False
            message = ""

            if not limited:
                message = f"Raw query executed."
            else:
                if is_aggregate:
                    message = f"Aggregate query completed."
                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}"
                    is_limited_result = True
                    message = f"Showing first {max_rows} rows."

            # Fetch results
            records = await conn.fetch(final_query)
            results = [dict(r) for r in records]
            
            if is_limited_result and len(results) < max_rows:
                is_limited_result = False
                message = f"Returned {len(results)} rows."
            elif is_limited_result:
                message = f"Showing first {len(results)} rows."

            columns = list(results[0].keys()) if results else []

            return {
                "success": True, 
                "results": jsonable_encoder(results, custom_encoder={uuid.UUID: str, ObjectId: str}), 
                "columns": columns, "rowCount": len(results), 
                "executionTime": time.time() - start_time,
                "is_aggregate": is_aggregate, "limited": is_limited_result, 
                "message": message, "error": None
            }
    except Exception as e:
        return {"success": False, "error": str(e), "executionTime": 0.0}

# ==============================================================================
# ROUTES
# ==============================================================================

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

@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_params = {
            "db": normalized_url, "q": query.sql_query, "l": limit_val, "lim": query.limited
        }
        # Direct Async Call (No Threadpool)
        result_dict = await coalescer.execute(
            "mysql", unique_params, _execute_async_mysql, 
            db_url=normalized_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 SqlQueryResponse(**final_result)
    except Exception as 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
        }
        # Direct Async Call (No Threadpool)
        result_dict = await coalescer.execute(
            "postgres", unique_params, _execute_async_postgres, 
            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:
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})

@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:
        parsed_query = parse_query_input(payload.generated_query)
        unique_params = {
            "uri": payload.mongo_uri, "db": payload.db_name, "col": payload.collection_name, 
            "q": parsed_query, "lim": payload.limited, "lrows": payload.limit_rows
        }
        # Mongo is Sync via PyMongo -> Uses run_in_threadpool
        result_data = await coalescer.execute(
            "mongo", unique_params, run_in_threadpool, 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_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 ---
class TestCalcRequest(BaseModel):
    value: int

def _heavy_calculation_task(value: int) -> int:
    time.sleep(0.1)
    return value * 2

@app.post("/api/test/parallel_calc_no_auth")
async def test_parallel_calc_endpoint(payload: TestCalcRequest):
    try:
        result = await run_in_threadpool(_heavy_calculation_task, value=payload.value)
        return {"success": True, "result": result, "process_id": os.getpid()}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/")
async def root():
    return {"message": "Code Execution Server is running"}

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

if __name__ == "__main__":
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 7860))
    num_workers = max(1, multiprocessing.cpu_count() - 2)
    print(f"Starting production server on {host}:{port} with {num_workers} workers...")
    uvicorn.run("controller:app", host=host, port=port, workers=num_workers, loop="asyncio")