from fastapi import APIRouter, HTTPException, UploadFile, File
from app.salary_analytics.core.state import state
from app.utils.logger import logger
import tempfile, os, time
from typing import Optional

router = APIRouter()



@router.post("/load-data")
async def load_data(source: str = "db", file: Optional[UploadFile] = File(None)):
    """
    Load data from either database or CSV file.
    
    Args:
        source (str): Source of data ('db' or 'csv')
        file (UploadFile, optional): CSV file to load (required if source is 'csv')
    
    Returns:
        dict: Status of data loading
    """
    start_time = time.time()
    try:
        if source not in ['db', 'csv']:
            logger.error(f"Invalid source: {source}")
            logger.info(f"Load data endpoint failed after {time.time() - start_time:.2f} seconds")
            raise HTTPException(status_code=400, detail="Source must be either 'db' or 'csv'")
        
        if source == 'csv' and not file:
            logger.error("No file provided for CSV source")
            logger.info(f"Load data endpoint failed after {time.time() - start_time:.2f} seconds")
            raise HTTPException(status_code=400, detail="File must be provided when loading from CSV")
        
        if source == 'csv':
            # Save uploaded file temporarily
            with tempfile.NamedTemporaryFile(delete=False, suffix='.csv') as temp_file:
                content = await file.read()
                temp_file.write(content)
                temp_file_path = temp_file.name
            
            try:
                success = state.pipeline.load_data(source='csv', file_path=temp_file_path)
            finally:
                # Clean up temporary file
                os.unlink(temp_file_path)
        else:
            success = state.pipeline.load_data(source='db')
        
        if not success:
            logger.error("Failed to load data")
            logger.info(f"Load data endpoint failed after {time.time() - start_time:.2f} seconds")
            raise HTTPException(status_code=500, detail="Failed to load data")
        
        response = {
            "status": "success",
            "message": f"Successfully loaded {len(state.pipeline.df)} rows of data",
            "columns": state.pipeline.df.columns.tolist(),
            "row_count": len(state.pipeline.df)
        }
        logger.info(f"Load data endpoint completed in {time.time() - start_time:.2f} seconds")
        return response
    except Exception as e:
        logger.error(f"Error loading data: {str(e)}")
        logger.info(f"Load data endpoint failed after {time.time() - start_time:.2f} seconds")
        raise HTTPException(status_code=500, detail=str(e))


@router.post("/load-data-with-file")
async def get_file_if_csv(source: str, file: Optional[UploadFile] = File(None)):
    """Dependency to handle file upload only when source is csv."""
    if source == 'csv' and not file:
        raise HTTPException(status_code=400, detail="File must be provided when loading from CSV")
    return file