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
import ssl
from typing import List, Optional, Dict, Any, TypeVar, Generic
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
from dotenv import load_dotenv
from pydantic import BaseModel

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

# --- Database Drivers (Sync - unused but kept for imports) ---
from bson import ObjectId
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 download_messages_log_helper import ChatLogRequest, ChatLogResponse, generate_chat_pdf
from extract_csv_metadata_service import extract_csv_metadata_logic
from mongo_service import convert_oid, execute_mongo_operation, extract_database_name, sanitize_json_input
from pdf_report_generation_helper import generate_and_upload_report
from pdf_report_generation_model import ReportGenerationRequest, ReportGenerationResponse
from pydantic_csv_analysis_model import AnalysisRequest, AnalysisResponse
from pydantic_csv_charts_model import ChartExecutionPayload, ChartExecutionResponse
from pydantic_migration_model import MigrationRequest, MigrationResponse
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, upload_file_to_supabase
from table_info_model import CsvFieldsRequest, CsvFieldsResponse
from file_uploader import ExternalFileUploader

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

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

uploader = ExternalFileUploader()

# ==============================================================================
# ASYNC CONNECTION POOL MANAGER (UPDATED WITH SSL FALLBACK)
# ==============================================================================
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 in self._pg_pools:
                return self._pg_pools[db_url]

            logger.info(f"Creating new AsyncPG pool for: {db_url[:25]}...")
            
            # Define connection logic
            async def attempt_connect(force_no_ssl: bool = False):
                # Common config
                config = {
                    "dsn": db_url,
                    "min_size": 1,
                    "max_size": 20,
                    "max_inactive_connection_lifetime": 300,
                    "command_timeout": 60
                }
                
                # If falling back, explicitly disable SSL (overrides DSN settings in most cases)
                if force_no_ssl:
                    config["ssl"] = False
                    
                return await asyncpg.create_pool(**config)

            try:
                # 1. Try Default Connection (uses DSN settings or defaults)
                pool = await attempt_connect(force_no_ssl=False)
                logger.info("✅ Connected to Postgres (Standard/SSL)")
                self._pg_pools[db_url] = pool
                
            except Exception as e:
                err_msg = str(e).lower()
                # 2. Check for SSL specific errors
                # "server does not support ssl" is the specific error from Postgres
                # "sslerror" covers generic python ssl issues
                if "ssl" in err_msg or "certificate" in err_msg or "connection reset" in err_msg:
                    logger.warning(f"⚠️ SSL Connection failed ({e}). Retrying with SSL DISABLED...")
                    try:
                        pool = await attempt_connect(force_no_ssl=True)
                        logger.info("✅ Connected to Postgres (No SSL Fallback)")
                        self._pg_pools[db_url] = pool
                    except Exception as retry_e:
                        logger.error(f"❌ Postgres Retry Failed: {retry_e}")
                        raise retry_e
                else:
                    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:
        # Check first without lock for speed
        if db_url in self._mysql_pools:
            return self._mysql_pools[db_url]

        async with self._lock:
            # Check again inside lock
            if db_url in self._mysql_pools:
                return self._mysql_pools[db_url]

            logger.info(f"Creating new AsyncMy pool for: {db_url[:25]}...")
            parsed = urlparse(db_url)
            qs = parse_qs(parsed.query)
            
            # Determine initial SSL settings
            initial_ssl_ctx = None
            if 'ssl-mode' in qs:
                mode = qs['ssl-mode'][0].upper()
                if mode != 'DISABLED':
                    initial_ssl_ctx = ssl.create_default_context()
                    initial_ssl_ctx.check_hostname = False
                    initial_ssl_ctx.verify_mode = ssl.CERT_NONE
            
            # Helper to create pool
            async def attempt_connect(ssl_context):
                return 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,
                    ssl=ssl_context 
                )

            try:
                # 1. Try with calculated SSL context
                pool = await attempt_connect(initial_ssl_ctx)
                logger.info("✅ Connected to MySQL")
                self._mysql_pools[db_url] = pool
            except Exception as e:
                # 2. If it fails and we were trying to use SSL, try disabling it
                if initial_ssl_ctx is not None:
                    logger.warning(f"⚠️ MySQL SSL handshake failed ({e}). Retrying without SSL...")
                    try:
                        pool = await attempt_connect(None) # Pass None to disable SSL
                        logger.info("✅ Connected to MySQL (No SSL Fallback)")
                        self._mysql_pools[db_url] = pool
                    except Exception as retry_e:
                        logger.error(f"❌ MySQL Retry Failed: {retry_e}")
                        raise retry_e
                else:
                    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()

async_pool_manager = AsyncPoolManager()

def get_dynamic_thread_limit():
    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):
    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
    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:
            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:
    """
    Normalizes the URI to ensure consistent cache keys.
    DOES NOT REMOVE parameters (like ssl-mode) blindly, but sorts them.
    """
    try:
        parsed_uri = urlparse(uri)
        query_params = parse_qs(parsed_uri.query)
        # Re-encode with sorting to ensure ?a=1&b=2 is same as ?b=2&a=1
        # We do NOT remove ssl-mode here anymore, so the pool manager knows about it
        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:
        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()
                
                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."

            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}
    
# ==============================================================================
# MIGRATION EXECUTORS
# ==============================================================================
async def _migrate_postgres(db_url: str, query: str) -> int:
    """Executes a SQL script (Create + Insert) in Postgres"""
    pool = await async_pool_manager.get_pg_pool(db_url)
    async with pool.acquire() as conn:
        async with conn.transaction():
            # asyncpg can execute scripts with multiple statements using .execute()
            # It returns string like "INSERT 0 100", "CREATE TABLE"
            await conn.execute(query)
            # We estimate rows by counting newlines in values or relying on frontend counts
            # For accurate counts, we'd need to parse the return tag, but 'execute' 
            # might run multiple statements.
            return 0 # Row counting is handled by frontend or specific return parsing

async def _migrate_mysql(db_url: str, query: str) -> int:
    """Executes a SQL script in MySQL"""
    pool = await async_pool_manager.get_mysql_pool(db_url)
    async with pool.acquire() as conn:
        async with conn.cursor() as cursor:
            # asyncmy usually expects single statements unless client flags allow multi
            # We will split by ';' for safety or execute raw if configured
            statements = [s.strip() for s in query.split(';') if s.strip()]
            for stmt in statements:
                await cursor.execute(stmt)
            return 0 

def _migrate_mongo(db_url: str, collection_name: str, data: List[Dict]) -> int:
    """Bulk Inserts documents into MongoDB"""
    # Use existing extract logic
    db_name = extract_database_name(db_url)
    if not db_name:
        raise ValueError("Could not determine database name from URI")
    
    from mongo_service import mongo_manager
    client = mongo_manager.get_client(db_url)
    db = client[db_name]
    collection = db[collection_name]
    
    if not data:
        return 0
        
    result = collection.insert_many(data)
    return len(result.inserted_ids)

# ==============================================================================
# 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
        }
        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
        }
        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:
        # 1. Determine Database Name
        final_db_name = payload.db_name
        
        # If DB Name is missing in payload, try to extract it from URI
        if not final_db_name:
            final_db_name = extract_database_name(payload.mongo_uri)

        # Final check
        if not final_db_name:
            raise HTTPException(
                status_code=400, 
                detail="Database name not provided in payload and could not be extracted from the URI path."
            )

        # 2. Parse Query & Determine Collection Name
        parsed_input = sanitize_json_input(payload.generated_query)
        final_query = parsed_input
        final_collection = payload.collection_name

        # Check for structured AI response: { "data": [ { "collectionName": "...", "pipeline": [...] } ] }
        if isinstance(parsed_input, dict) and "data" in parsed_input and isinstance(parsed_input["data"], list):
            if len(parsed_input["data"]) > 0:
                extracted_data = parsed_input["data"][0]
                
                # Extract Pipeline
                if "pipeline" in extracted_data:
                    final_query = extracted_data["pipeline"]
                
                # Extract Collection Name (Overrides payload if present)
                if "collectionName" in extracted_data and extracted_data["collectionName"]:
                    final_collection = extracted_data["collectionName"]

        if not final_collection:
             raise HTTPException(status_code=400, detail="Collection name could not be determined from payload or query.")

        # 3. Execute
        unique_params = {
            "uri": payload.mongo_uri, 
            "db": final_db_name, 
            "col": final_collection,
            "q": str(final_query),
            "lim": payload.limited
        }
        
        final_query = convert_oid(final_query)

        result_data = await coalescer.execute(
            "mongo", 
            unique_params, 
            run_in_threadpool, 
            execute_mongo_operation, 
            mongo_uri=payload.mongo_uri, 
            db_name=final_db_name, 
            collection_name=final_collection, 
            query=final_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 HTTPException as he:
        raise he
    except Exception as e:
        logger.error(f"Endpoint Error: {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:
            unique_name = f"{uuid.uuid4()}.png"
            base64_str = base64.b64encode(image_bytes).decode('utf-8')
            file_url = await run_in_threadpool(uploader.upload_file, file_bytes=image_bytes, file_name=unique_name)
            if file_url:
                return ChartExecutionResponse(status="success", image_url=file_url, output_log=logs, request_id=request_id)
            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})
    
# --- Migration Route ---
@app.post("/api/migrate", response_model=MigrationResponse)
async def execute_migration_endpoint(payload: MigrationRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    start_time = time.time()
    
    try:
        rows_processed = 0
        
        # 1. PostgreSQL Migration
        if payload.db_type == 'postgresql':
            if not payload.migration_query:
                raise ValueError("Migration query is required for SQL databases")
            
            clean_url = normalize_postgres_uri(payload.db_url)
            await _migrate_postgres(db_url=clean_url, query=payload.migration_query)
            # Row count is difficult to get exactly from a multi-statement script 
            # without complex parsing, usually the frontend knows how much data it sent.
            rows_processed = payload.migration_query.upper().count("VALUES") # Rough estimate or 0
            
        # 2. MySQL Migration
        elif payload.db_type == 'mysql':
             if not payload.migration_query:
                raise ValueError("Migration query is required for SQL databases")
                
             clean_url = normalize_mysql_uri(payload.db_url)
             # Note: run_in_threadpool might be needed if async driver issues arise, 
             # but here we defined async wrappers, so we await directly or use logic
             # Since _migrate_mysql is async, we await it.
             await _migrate_mysql(db_url=clean_url, query=payload.migration_query)
             rows_processed = payload.migration_query.upper().count("VALUES")

        # 3. MongoDB Migration
        elif payload.db_type == 'mongodb':
            if not payload.migration_data:
                 raise ValueError("Migration data (JSON) is required for MongoDB")
            if not payload.collection_name:
                 raise ValueError("Collection name is required for MongoDB")
                 
            # Mongo operations are blocking in pymongo, so run in threadpool
            rows_processed = await run_in_threadpool(
                _migrate_mongo, 
                db_url=payload.db_url, 
                collection_name=payload.collection_name, 
                data=payload.migration_data
            )
            
        else:
            raise ValueError(f"Unsupported database type: {payload.db_type}")

        duration = f"{time.time() - start_time:.2f}s"
        
        logger.info(f"Migration Complete: {rows_processed} batches processed in {duration}")
        
        return MigrationResponse(
            success=True,
            rowsProcessed=rows_processed,
            duration=duration,
            errors=[],
            request_id=request_id
        )

    except Exception as e:
        logger.error(f"Migration Failed: {e}")
        duration = f"{time.time() - start_time:.2f}s"
        return MigrationResponse(
            success=False,
            rowsProcessed=0,
            duration=duration,
            errors=[str(e)],
            request_id=request_id
        )
        
# --- CSV Metadata ---
@app.post("/api/get_csv_fields", response_model=CsvFieldsResponse)
async def get_csv_fields_endpoint(payload: CsvFieldsRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    try:
        fields = await run_in_threadpool(extract_csv_metadata_logic, csv_url=payload.csv_url)
        return CsvFieldsResponse(success=True, fields=fields, request_id=request_id)
    except Exception as e:
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})
        
# --- Chat Messages Log ---
@app.post("/api/get_chat_log", response_model=ChatLogResponse)
async def get_chat_log_endpoint(payload: ChatLogRequest, token: str = Depends(validate_token)):
    request_id = str(uuid.uuid4())[:8]
    generated_pdf_path = None
    
    try:
        # 1. Generate the PDF locally
        # run_in_threadpool is crucial here because PDF generation is CPU/IO blocking
        generated_pdf_path = await run_in_threadpool(
            generate_chat_pdf,
            messages=payload.messages,
            title="Chat Conversation Transcript",
            chat_id=payload.chatId,
            output_dir="chat_pdfs", 
            include_stats=True,
        )

        if not generated_pdf_path or not os.path.exists(generated_pdf_path):
            raise Exception("PDF generation failed or file not found.")

        # 2. Upload to Supabase
        # Extract filename from path (e.g., "chat_log_123.pdf")
        filename = os.path.basename(generated_pdf_path)
        
        public_url = await run_in_threadpool(
            upload_file_to_supabase,
            file_path=generated_pdf_path,
            file_name=filename,
            chat_id=payload.chatId
        )

        return ChatLogResponse(
            success=True,
            pdf_url=public_url,
            request_id=request_id,
            message="PDF generated and uploaded successfully."
        )

    except Exception as e:
        logger.error(f"Chat Log Error [{request_id}]: {str(e)}")
        raise HTTPException(
            status_code=500, 
            detail={"success": False, "error": str(e), "request_id": request_id}
        )
        
    finally:
        # 3. Cleanup: Delete the local file to save space
        if generated_pdf_path and os.path.exists(generated_pdf_path):
            try:
                os.remove(generated_pdf_path)
                logger.info(f"Cleaned up local PDF: {generated_pdf_path}")
            except Exception as cleanup_error:
                logger.warning(f"Failed to delete local PDF {generated_pdf_path}: {cleanup_error}")

# --- PDF Report Generation ---
@app.post("/api/generate_pdf_report", response_model=ReportGenerationResponse)
async def generate_pdf_report_endpoint(
    payload: ReportGenerationRequest, 
    token: str = Depends(validate_token)
):
    """
    Generates a PDF report, uploads it to Supabase, and returns the public URL.
    
    This endpoint:
    1. Validates the request payload
    2. Converts Pydantic models to dictionaries
    3. Generates the PDF with proper formatting
    4. Uploads to Supabase storage
    5. Cleans up temporary files
    6. Returns the public URL
    """
    request_id = str(uuid.uuid4())[:8]
    
    try:
        logger.info(f"[{request_id}] Starting PDF report generation for chat: {payload.chat_id}")
        
        # Convert Pydantic models to dictionaries for the generator
        # Use .model_dump() for Pydantic v2, .dict() for v1
        try:
            config_content = [section.model_dump() for section in payload.sections]
        except AttributeError:
            # Fallback for Pydantic v1
            config_content = [section.dict() for section in payload.sections]
        
        logger.info(f"[{request_id}] Processing {len(config_content)} sections")
        
        # Call the async helper function
        public_url = await generate_and_upload_report(
            config_content=config_content,
            file_name=payload.file_name,
            chat_id=payload.chat_id,
            title=payload.title,
            subtitle=payload.subtitle,
            author=payload.author,
            department=payload.department,
            output_dir="temp_reports",
            confidential=payload.confidential
        )
        
        # Prepare final filename
        final_filename = payload.file_name
        if not final_filename.endswith('.pdf'):
            final_filename += '.pdf'
        
        logger.info(f"[{request_id}] Report generated successfully: {public_url}")
        
        return ReportGenerationResponse(
            success=True,
            pdf_report_file_url=public_url,
            file_name=final_filename,
            request_id=request_id
        )
        
    except Exception as e:
        logger.error(f"[{request_id}] Report generation failed: {e}", exc_info=True)
        return ReportGenerationResponse(
            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")