controller.py

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

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

# --- Database Drivers ---
from bson import ObjectId
import mysql.connector
import psycopg2
from psycopg2.extras import RealDictCursor

# --- 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
    # Shutdown logic (if any) goes here

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

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

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

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

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

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

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

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

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

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


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

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

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


# ==============================================================================
# SHARED HELPER FUNCTIONS
# ==============================================================================

def is_aggregate_query(query: str) -> bool:
    """Checks for aggregate keywords."""
    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

# ==============================================================================
# MYSQL LOGIC
# ==============================================================================

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 _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:
        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
        }
        connection = mysql.connector.connect(**db_config)
        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 and connection.is_connected(): connection.close()

# ==============================================================================
# POSTGRES LOGIC
# ==============================================================================

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_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:
        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}
        connection = psycopg2.connect(**db_config)
        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: connection.close()


# ==============================================================================
# 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:
        parsed_query = parse_query_input(payload.generated_query)
        result_data = await 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:
        # 1. Execute Code
        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)

        # 2. Handle Output Format
        if payload.return_base64:
            # OPTION A: Return Base64 (No Supabase Upload)
            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:
            # OPTION B: Upload to Supabase (Standard behavior)
            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
        result_dict = await run_in_threadpool(_run_mysql_synchronously, db_url=normalized_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited)
        result_dict["request_id"] = request_id
        return SqlQueryResponse(**result_dict)
    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
        result_dict = await run_in_threadpool(_run_postgres_synchronously, db_url=clean_url, sql_query=query.sql_query, max_rows=limit_val, limited=query.limited)
        result_dict["request_id"] = request_id
        return PgQueryResponse(**result_dict)
    except Exception as e:
        raise HTTPException(status_code=500, detail={"success": False, "error": str(e), "request_id": request_id})

@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})

# ==============================================================================
# NEW: BATCH HANDLER LOGIC (Scaling for 1000+ Requests)
# ==============================================================================

async def batch_parallel_handler(func, requests: List[Any], token: str):
    """
    Executes multiple requests in parallel within a single worker process 
    using asyncio.gather, utilizing the high-token thread pool.
    """
    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)

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

if __name__ == "__main__":
    import uvicorn
    host = os.getenv("HOST", "0.0.0.0")
    port = int(os.getenv("PORT", 8000))
    
    # Scale processes: 16 Cores - 2 = 14 Workers
    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="auto",
    )