upload.py

"""Upload router for the Memory Tracker API."""

from fastapi import APIRouter, Depends, HTTPException
from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy import select, desc, func, update, and_, delete
from sqlalchemy.exc import IntegrityError
from datetime import datetime
import logging

from .. import schemas, crud, models
from ..database import get_database
from ..auth import get_current_token

router = APIRouter(prefix="/api", tags=["upload"])


async def cleanup_old_flamegraphs_if_needed(
    db: AsyncSession, environment_id: str, max_flamegraphs: int = 100
):
    """
    Set flamegraph_html to NULL for runs older than the last max_flamegraphs in this environment.
    """
    logger = logging.getLogger(__name__)

    keep_runs_result = await db.execute(
        select(models.Run.run_id)
        .where(models.Run.environment_id == environment_id)
        .order_by(desc(models.Run.timestamp))
        .limit(max_flamegraphs)
    )
    keep_run_ids = [row[0] for row in keep_runs_result.fetchall()]

    if not keep_run_ids:
        return 0

    total_flamegraphs_result = await db.execute(
        select(func.count())
        .select_from(models.BenchmarkResult)
        .join(models.Run)
        .where(
            models.Run.environment_id == environment_id,
            models.BenchmarkResult.flamegraph_html.is_not(None),
        )
    )
    total_flamegraphs_before = total_flamegraphs_result.scalar()

    count_query = (
        select(func.count())
        .select_from(models.BenchmarkResult)
        .where(
            and_(
                models.BenchmarkResult.run_id.not_in(keep_run_ids),
                models.BenchmarkResult.flamegraph_html.is_not(None),
            )
        )
    )
    count_result = await db.execute(count_query)
    rows_to_clean = count_result.scalar()

    if rows_to_clean == 0:
        logger.info(
            f"No flamegraphs need cleaning for environment '{environment_id}' (all {total_flamegraphs_before} flamegraphs are from recent runs)"
        )
        return 0

    cleanup_query = (
        update(models.BenchmarkResult)
        .where(
            and_(
                models.BenchmarkResult.run_id.not_in(keep_run_ids),
                models.BenchmarkResult.flamegraph_html.is_not(None),
            )
        )
        .values(flamegraph_html=None)
    )

    try:
        await db.execute(cleanup_query)
        await db.commit()

        verify_result = await db.execute(count_query)
        remaining = verify_result.scalar()
        actual_cleaned = rows_to_clean - remaining

        final_flamegraphs_result = await db.execute(
            select(func.count())
            .select_from(models.BenchmarkResult)
            .join(models.Run)
            .where(
                models.Run.environment_id == environment_id,
                models.BenchmarkResult.flamegraph_html.is_not(None),
            )
        )
        total_flamegraphs_after = final_flamegraphs_result.scalar()

        if actual_cleaned > 0:
            logger.info(
                f"Cleaned up {actual_cleaned} flamegraphs for environment '{environment_id}': {total_flamegraphs_before} → {total_flamegraphs_after} flamegraphs remaining"
            )
        else:
            logger.error(
                f"Cleanup FAILED for environment '{environment_id}'. Expected to clean {rows_to_clean}, but {remaining} still remain"
            )

        return actual_cleaned

    except Exception as e:
        logger.error(
            f"Exception during cleanup for environment '{environment_id}': {e}"
        )
        await db.rollback()
        raise


@router.post("/upload-run", response_model=dict)
async def upload_worker_run(
    upload_data: schemas.WorkerRunUpload,
    db: AsyncSession = Depends(get_database),
    current_token: models.AuthToken = Depends(get_current_token),
):
    logger = logging.getLogger(__name__)

    logger.info(
        f"Authenticated upload request from token '{current_token.name}' for binary_id='{upload_data.binary_id}', environment_id='{upload_data.environment_id}'"
    )
    logger.debug(
        f"Upload contains {len(upload_data.benchmark_results)} benchmark results"
    )

    metadata = upload_data.metadata

    # Extract commit information
    commit_info = metadata.get("commit", {})
    commit_sha = commit_info.get("hexsha")
    if not commit_sha:
        logger.error("Upload failed: Missing commit SHA in metadata")
        raise HTTPException(status_code=400, detail="Missing commit SHA in metadata")

    logger.info(f"Processing upload for commit {commit_sha[:8]}")

    # Extract Python version
    version_info = metadata.get("version", {})
    python_version = schemas.PythonVersion(
        major=version_info.get("major"),
        minor=version_info.get("minor"),
        patch=version_info.get(
            "micro", 0
        ),  # 'micro' in metadata maps to 'patch' in our schema
    )
    logger.debug(
        f"Extracted Python version: {python_version.major}.{python_version.minor}.{python_version.patch}"
    )

    # Use provided binary_id and environment_id from worker
    binary_id = upload_data.binary_id
    environment_id = upload_data.environment_id

    # Validate binary exists
    logger.debug(f"Validating binary '{binary_id}' exists")
    binary = await crud.get_binary_by_id(db, binary_id=binary_id)
    if not binary:
        logger.error(f"Upload failed: Binary '{binary_id}' not found")
        raise HTTPException(
            status_code=400,
            detail=f"Binary '{binary_id}' not found. Binaries must be pre-registered.",
        )
    logger.info(f"Binary '{binary_id}' validated successfully")

    # Validate environment exists
    logger.debug(f"Validating environment '{environment_id}' exists")
    environment = await crud.get_environment_by_id(db, environment_id=environment_id)
    if not environment:
        logger.error(f"Upload failed: Environment '{environment_id}' not found")
        raise HTTPException(
            status_code=400,
            detail=f"Environment '{environment_id}' not found. Environments must be pre-registered.",
        )
    logger.info(f"Environment '{environment_id}' validated successfully")

    # Validate configure flags - the registered binary flags must be a subset of uploaded flags
    configure_vars = metadata.get("configure_vars", {})
    uploaded_config_args = configure_vars.get("CONFIG_ARGS", "")
    
    # Strip quotes from each flag and filter out non-configure flags
    def clean_flag(flag):
        """Remove surrounding quotes from a flag."""
        return flag.strip().strip("'\"")
    
    # Split and clean uploaded flags, filtering out environment variables
    raw_uploaded_flags = uploaded_config_args.split() if uploaded_config_args else []
    uploaded_flags = set()
    for flag in raw_uploaded_flags:
        cleaned = clean_flag(flag)
        # Only include flags that start with -- (configure flags)
        if cleaned.startswith('--'):
            uploaded_flags.add(cleaned)
    
    # Clean registered flags too for consistency
    registered_flags = {clean_flag(flag) for flag in (binary.flags or [])}

    logger.debug(
        f"Configure flags validation: registered={sorted(registered_flags)}, uploaded={sorted(uploaded_flags)}"
    )
    logger.debug(f"Raw CONFIG_ARGS from metadata: '{uploaded_config_args}'")

    # Check if registered flags are a subset of uploaded flags
    if registered_flags and not registered_flags.issubset(uploaded_flags):
        missing_flags = registered_flags - uploaded_flags
        logger.error(
            f"Upload failed: Configure flags mismatch for binary '{binary_id}'. "
            f"Missing flags: {sorted(missing_flags)}, "
            f"Required: {sorted(registered_flags)}, "
            f"Provided: {sorted(uploaded_flags)}, "
            f"Raw CONFIG_ARGS: '{uploaded_config_args}'"
        )
        raise HTTPException(
            status_code=400,
            detail=f"Binary '{binary_id}' requires configure flags {sorted(list(registered_flags))} but upload only has {sorted(list(uploaded_flags))}. "
            f"Registered configure flags must be a subset of upload configure flags.",
        )
    logger.info(f"Configure flags validation passed for binary '{binary_id}'")

    # Note: Duplicate commits are now prevented by database unique constraint on (commit_sha, binary_id, environment_id)

    # Create or get commit
    logger.debug(f"Looking up commit {commit_sha[:8]} in database")
    commit = await crud.get_commit_by_sha(db, sha=commit_sha)
    if not commit:
        logger.info(f"Commit {commit_sha[:8]} not found, creating new commit record")
        # Create commit from metadata
        commit_data = schemas.CommitCreate(
            sha=commit_sha,
            timestamp=datetime.fromisoformat(
                commit_info.get("committed_date", "").replace("Z", "+00:00")
            ),
            message=commit_info.get("message", ""),
            author=commit_info.get("author", ""),
            python_version=python_version,
        )
        try:
            commit = await crud.create_commit(db, commit_data)
            logger.info(f"Successfully created commit record for {commit_sha[:8]}")
        except Exception as e:
            logger.error(f"Failed to create commit record for {commit_sha[:8]}: {e}")
            raise HTTPException(
                status_code=500, detail=f"Failed to create commit record: {str(e)}"
            )
    else:
        logger.debug(f"Found existing commit record for {commit_sha[:8]}")

    # Create run
    run_id = f"run_{commit_sha[:8]}_{binary_id}_{environment_id}_{int(datetime.now().timestamp())}"
    logger.info(f"Creating run with ID: {run_id}")
    run_data = schemas.RunCreate(
        run_id=run_id,
        commit_sha=commit_sha,
        binary_id=binary_id,
        environment_id=environment_id,
        python_version=python_version,
        timestamp=datetime.now(),
    )

    try:
        await crud.create_run(db, run_data)
        logger.info(f"Successfully created run record: {run_id}")

        # Create benchmark results
        created_results = []
        logger.info(
            f"Processing {len(upload_data.benchmark_results)} benchmark results"
        )
        for i, benchmark_result in enumerate(upload_data.benchmark_results, 1):
            logger.debug(
                f"Processing benchmark result {i}/{len(upload_data.benchmark_results)}: {benchmark_result.benchmark_name}"
            )
            # Convert worker format to internal format
            stats_json = benchmark_result.stats_json

            # Extract key metrics from the stats JSON
            result_json = schemas.BenchmarkResultJson(
                high_watermark_bytes=stats_json.get("metadata", {}).get(
                    "peak_memory", 0
                ),
                allocation_histogram=[
                    [item["min_bytes"], item["count"]]
                    for item in stats_json.get("allocation_size_histogram", [])
                ],
                total_allocated_bytes=stats_json.get("total_bytes_allocated", 0),
                top_allocating_functions=[
                    schemas.TopAllocatingFunction(
                        function=alloc["location"],
                        count=alloc.get("count", 0),
                        total_size=alloc["size"],
                    )
                    for alloc in stats_json.get("top_allocations_by_size", [])[:10]
                ],
                benchmark_name=benchmark_result.benchmark_name,
            )

            result_data = schemas.BenchmarkResultCreate(
                run_id=run_id,
                benchmark_name=benchmark_result.benchmark_name,
                result_json=result_json,
                flamegraph_html=benchmark_result.flamegraph_html,
            )
            try:
                created_result = await crud.create_benchmark_result(db, result_data)
                created_results.append(created_result.id)
                logger.debug(
                    f"Successfully created benchmark result for {benchmark_result.benchmark_name}"
                )
            except Exception as e:
                logger.error(
                    f"Failed to create benchmark result for {benchmark_result.benchmark_name}: {e}"
                )
                raise

        logger.info(
            f"Upload completed successfully: run_id={run_id}, created {len(created_results)} benchmark results"
        )

        # Clean up old flamegraphs if we have more than 100 runs for this environment
        await cleanup_old_flamegraphs_if_needed(db, environment_id, max_flamegraphs=100)

        # Clear any memray build failures for this binary+environment since upload succeeded
        await db.execute(
            delete(models.MemrayBuildFailure).where(
                models.MemrayBuildFailure.binary_id == binary_id,
                models.MemrayBuildFailure.environment_id == environment_id
            )
        )

        return {
            "message": "Worker run uploaded successfully",
            "run_id": run_id,
            "commit_sha": commit_sha,
            "binary_id": binary_id,
            "environment_id": environment_id,
            "results_created": len(created_results),
            "result_ids": created_results,
        }

    except IntegrityError as e:
        # The only unique constraint that can fire here is
        # unique_commit_binary_env on the runs table.
        error_str = str(e).lower()
        if "unique_commit_binary_env" in error_str or (
            "commit_sha" in error_str
            and "binary_id" in error_str
            and "environment_id" in error_str
        ):
            logger.error(
                f"Upload failed: Duplicate run for commit {commit_sha[:8]}, binary '{binary_id}', environment '{environment_id}'"
            )
            raise HTTPException(
                status_code=409,
                detail=f"A run already exists for commit {commit_sha[:8]} with binary '{binary_id}' and environment '{environment_id}'. Duplicate uploads are not allowed.",
            )
        else:
            logger.error(f"Database integrity error during upload: {e}")
            raise HTTPException(
                status_code=500, detail=f"Database integrity error: {str(e)}"
            )
    except HTTPException:
        # Re-raise HTTP exceptions (validation errors) as-is
        raise
    except Exception as e:
        logger.error(f"Unexpected error during upload processing: {e}")
        raise HTTPException(
            status_code=500, detail=f"Failed to upload worker run: {str(e)}"
        )


@router.post("/report-memray-failure", response_model=dict)
async def report_memray_failure(
    failure_data: schemas.MemrayFailureReport,
    db: AsyncSession = Depends(get_database),
    current_token: models.AuthToken = Depends(get_current_token),
):
    """Report a memray build failure for a specific commit and environment."""
    logger = logging.getLogger(__name__)
    
    logger.info(
        f"Authenticated memray failure report from token '{current_token.name}' for commit {failure_data.commit_sha[:8]}, "
        f"binary_id='{failure_data.binary_id}', environment_id='{failure_data.environment_id}'"
    )
    
    # Validate binary exists
    binary = await crud.get_binary_by_id(db, binary_id=failure_data.binary_id)
    if not binary:
        logger.error(f"Memray failure report failed: Binary '{failure_data.binary_id}' not found")
        raise HTTPException(
            status_code=400,
            detail=f"Binary '{failure_data.binary_id}' not found."
        )
    
    # Validate environment exists
    environment = await crud.get_environment_by_id(db, environment_id=failure_data.environment_id)
    if not environment:
        logger.error(f"Memray failure report failed: Environment '{failure_data.environment_id}' not found")
        raise HTTPException(
            status_code=400,
            detail=f"Environment '{failure_data.environment_id}' not found."
        )
    
    # Check if this failure is newer than any existing failure for this binary+environment
    existing_failure = await db.execute(
        select(models.MemrayBuildFailure)
        .where(
            models.MemrayBuildFailure.binary_id == failure_data.binary_id,
            models.MemrayBuildFailure.environment_id == failure_data.environment_id
        )
    )
    existing_failure = existing_failure.scalars().first()
    
    if existing_failure:
        # Check if the new failure is from a newer commit
        if failure_data.commit_timestamp <= existing_failure.commit_timestamp:
            logger.info(
                f"Ignoring memray failure for {failure_data.commit_sha[:8]} as it's not newer than existing failure"
            )
            return {
                "message": "Memray failure ignored (not newer than existing failure)",
                "commit_sha": failure_data.commit_sha,
                "binary_id": failure_data.binary_id,
                "environment_id": failure_data.environment_id,
            }
        
        # Update existing failure with newer information
        existing_failure.commit_sha = failure_data.commit_sha
        existing_failure.error_message = failure_data.error_message
        existing_failure.failure_timestamp = datetime.now()
        existing_failure.commit_timestamp = failure_data.commit_timestamp
        
        try:
            await db.commit()
            logger.info(f"Updated existing memray failure record for binary '{failure_data.binary_id}', environment '{failure_data.environment_id}'")
        except Exception as e:
            logger.error(f"Failed to update memray failure record: {e}")
            raise HTTPException(
                status_code=500, detail=f"Failed to update memray failure record: {str(e)}"
            )
    else:
        # Create new failure record
        failure_record = models.MemrayBuildFailure(
            commit_sha=failure_data.commit_sha,
            binary_id=failure_data.binary_id,
            environment_id=failure_data.environment_id,
            error_message=failure_data.error_message,
            failure_timestamp=datetime.now(),
            commit_timestamp=failure_data.commit_timestamp
        )
        
        try:
            db.add(failure_record)
            await db.commit()
            logger.info(f"Created new memray failure record for commit {failure_data.commit_sha[:8]}")
        except Exception as e:
            logger.error(f"Failed to create memray failure record: {e}")
            raise HTTPException(
                status_code=500, detail=f"Failed to create memray failure record: {str(e)}"
            )
    
    return {
        "message": "Memray failure reported successfully",
        "commit_sha": failure_data.commit_sha,
        "binary_id": failure_data.binary_id,
        "environment_id": failure_data.environment_id,
    }