"""Direct file staging for small files.

This module provides a fast path for staging small files directly to DuckDB,
bypassing Dagster job overhead while still reporting AssetMaterializations
for observability in the Dagster UI.

For small files (< SMALL_FILE_STAGING_THRESHOLD_MB), this approach:
- Stages files immediately during the HTTP request
- Reports an AssetMaterialization to Dagster for observability
- Returns control to the caller in milliseconds instead of seconds

For large files, the regular Dagster job should be used for:
- Async processing that won't timeout
- Progress streaming for long operations
- Retry logic for unreliable operations
"""

import time
from typing import Any

from sqlalchemy.orm import Session

from robosystems.config import env
from robosystems.logger import logger


async def stage_file_directly(
  db: Session,
  file_id: str,
  graph_id: str,
  table_id: str,
  s3_key: str,
  file_size_bytes: int,
  row_count: int | None = None,
) -> dict[str, Any]:
  """
  Stage a file directly to DuckDB without Dagster orchestration.

  This is the fast path for small files. It:
  1. Gets all uploaded files for the table
  2. Calls the Graph API to create/update the DuckDB staging table
  3. Marks the file as staged in the database
  4. Reports an AssetMaterialization to Dagster for observability

  Args:
      db: SQLAlchemy database session (from FastAPI dependency injection)
      file_id: The file ID to stage
      graph_id: The graph database ID
      table_id: The table ID
      s3_key: The S3 key of the file
      file_size_bytes: Size of the file in bytes
      row_count: Optional row count for the file

  Returns:
      Dict with staging result including status, duration, etc.
  """
  # Lazy imports to avoid circular dependencies
  from robosystems.graph_api.client.factory import GraphClientFactory
  from robosystems.models.core import GraphFile, GraphTable

  start_time = time.time()

  logger.info(
    f"Direct staging file {file_id} to graph {graph_id} (fast path for small file)"
  )

  try:
    # Get the file and table
    graph_file = GraphFile.get_by_id(file_id, db)
    if not graph_file:
      return {
        "status": "error",
        "message": f"File {file_id} not found",
        "file_id": file_id,
      }

    table = GraphTable.get_by_id(table_id, db)
    if not table:
      return {
        "status": "error",
        "message": f"Table {table_id} not found",
        "file_id": file_id,
      }

    # Get all uploaded files for this table (for proper staging)
    all_files = GraphFile.get_all_for_table(table_id, db)
    uploaded_files = [f for f in all_files if f.upload_status == "uploaded"]

    if not uploaded_files:
      return {
        "status": "skipped",
        "message": "No uploaded files to stage",
        "file_id": file_id,
      }

    # Build file list with S3 URIs
    bucket = env.USER_DATA_BUCKET
    s3_files = [f"s3://{bucket}/{f.s3_key}" for f in uploaded_files]
    file_id_map = {f"s3://{bucket}/{f.s3_key}": f.id for f in uploaded_files}

    logger.info(
      f"Direct staging {len(s3_files)} files to DuckDB table {table.table_name}"
    )

    # Stage via Graph API
    client = await GraphClientFactory.create_client(
      graph_id=graph_id, operation_type="write"
    )

    try:
      # Check if the DuckDB table already exists to use incremental INSERT INTO
      existing_tables = await client.list_tables(graph_id)
      existing_table_names = [t["table_name"] for t in existing_tables]

      if table.table_name in existing_table_names:
        # Table exists - use INSERT INTO for just the new file (incremental)
        new_file_s3 = f"s3://{bucket}/{s3_key}"
        logger.info(
          f"Table {table.table_name} exists, using INSERT INTO for file {s3_key}"
        )
        staging_result = await client.insert_into_table(
          graph_id=graph_id,
          table_name=table.table_name,
          s3_pattern=[new_file_s3],
          deduplicate=True,
          file_id_map={new_file_s3: file_id},
        )
      else:
        # Table does not exist - create with all files (first-file path)
        staging_result = await client.create_table(
          graph_id=graph_id,
          table_name=table.table_name,
          s3_pattern=s3_files,
          file_id_map=file_id_map,
        )
      logger.debug(f"DuckDB staging result: {staging_result}")
    finally:
      await client.close()

    # Mark file as staged
    graph_file.mark_duckdb_staged(session=db, row_count=row_count or 0)

    duration_ms = (time.time() - start_time) * 1000

    logger.info(f"Direct staging completed for file {file_id} in {duration_ms:.2f}ms")

    # Report AssetMaterialization to Dagster for observability (fire-and-forget with timeout)
    from robosystems.dagster.reporting import report_asset_materialization

    await report_asset_materialization(
      asset_key="user_graph_file_staging",
      description=f"Direct staging of {len(s3_files)} file(s) to table {table.table_name}",
      metadata={
        "file_id": file_id,
        "graph_id": graph_id,
        "table_name": table.table_name,
        "file_size_bytes": file_size_bytes,
        "row_count": row_count or 0,
        "duration_ms": duration_ms,
        "files_staged": len(s3_files),
        "staging_method": "direct",
      },
    )

    return {
      "status": "success",
      "file_id": file_id,
      "graph_id": graph_id,
      "table_name": table.table_name,
      "files_staged": len(s3_files),
      "duckdb_status": "staged",
      "duration_ms": duration_ms,
      "method": "direct",
    }

  except Exception as e:
    duration_ms = (time.time() - start_time) * 1000
    logger.error(f"Direct staging failed for file {file_id}: {e}")
    return {
      "status": "error",
      "file_id": file_id,
      "message": str(e),
      "duration_ms": duration_ms,
      "method": "direct",
    }
