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
from typing import List, Optional, Dict, Any, TypeVar, Generic

# --- 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

# --- Database Drivers ---
from bson import ObjectId
from psycopg2.extras import RealDictCursor
from urllib.parse import urlparse, parse_qs, urlencode, urlunparse

# --- Database Pool ---
from mysql.connector import pooling
from psycopg2 import pool as pg_pool
from threading import Lock

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

@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup logic
    to_thread.current_default_thread_limiter().total_tokens = 2000
    logger.info("Worker Process Started: Thread pool capacity set to 2000.")
    yield

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):
        if db_url in self._pg_pools and conn:
            self._pg_pools[db_url].putconn(conn)

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": ""}

    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:
        response["error"] = str(e)
        return response
    finally:
        if cursor: cursor.close()
        if connection: connection.close() # Returns to pool

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": ""}
    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:
        if connection: connection.rollback()
        response["error"] = str(e)
        return response
    finally:
        if cursor: cursor.close()
        if connection: pool_manager.return_postgres_connection(db_url, connection)


# ==============================================================================
# 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",
    )