# [DEF:AppModule:Module]
# @SEMANTICS: app, main, entrypoint, fastapi
# @PURPOSE: The main entry point for the FastAPI application. It initializes the app, configures CORS, sets up dependencies, includes API routers, and defines the WebSocket endpoint for log streaming.
# @LAYER: UI (API)
# @RELATION: Depends on the dependency module and API route modules.
import sys
from pathlib import Path

# project_root is used for static files mounting
project_root = Path(__file__).resolve().parent.parent.parent

from fastapi import FastAPI, WebSocket, WebSocketDisconnect, Depends, Request, HTTPException
from starlette.middleware.sessions import SessionMiddleware
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse
import asyncio
import os

from .dependencies import get_task_manager, get_scheduler_service
from .core.utils.network import NetworkError
from .core.logger import logger, belief_scope
from .api.routes import plugins, tasks, settings, environments, mappings, migration, connections, git, storage, admin, llm
from .api import auth
from .core.database import init_db

# [DEF:App:Global]
# @SEMANTICS: app, fastapi, instance
# @PURPOSE: The global FastAPI application instance.
app = FastAPI(
    title="Superset Tools API",
    description="API for managing Superset automation tools and plugins.",
    version="1.0.0",
)
# [/DEF:App:Global]

# [DEF:startup_event:Function]
# @PURPOSE: Handles application startup tasks, such as starting the scheduler.
# @PRE: None.
# @POST: Scheduler is started.
# Startup event
@app.on_event("startup")
async def startup_event():
    with belief_scope("startup_event"):
        scheduler = get_scheduler_service()
    scheduler.start()
# [/DEF:startup_event:Function]

# [DEF:shutdown_event:Function]
# @PURPOSE: Handles application shutdown tasks, such as stopping the scheduler.
# @PRE: None.
# @POST: Scheduler is stopped.
# Shutdown event
@app.on_event("shutdown")
async def shutdown_event():
    with belief_scope("shutdown_event"):
        scheduler = get_scheduler_service()
    scheduler.stop()
# [/DEF:shutdown_event:Function]

# Configure Session Middleware (required by Authlib for OAuth2 flow)
from .core.auth.config import auth_config
app.add_middleware(SessionMiddleware, secret_key=auth_config.SECRET_KEY)

# Configure CORS
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust this in production
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# [DEF:log_requests:Function]
# @PURPOSE: Middleware to log incoming HTTP requests and their response status.
# @PRE: request is a FastAPI Request object.
# @POST: Logs request and response details.
# @PARAM: request (Request) - The incoming request object.
# @PARAM: call_next (Callable) - The next middleware or route handler.
@app.exception_handler(NetworkError)
async def network_error_handler(request: Request, exc: NetworkError):
    with belief_scope("network_error_handler"):
        logger.error(f"Network error: {exc}")
        return HTTPException(
            status_code=503,
            detail="Environment unavailable. Please check if the Superset instance is running."
        )

@app.middleware("http")
async def log_requests(request: Request, call_next):
    # Avoid spamming logs for polling endpoints
    is_polling = request.url.path.endswith("/api/tasks") and request.method == "GET"
    
    if not is_polling:
        logger.info(f"Incoming request: {request.method} {request.url.path}")
    
    try:
        response = await call_next(request)
        if not is_polling:
            logger.info(f"Response status: {response.status_code} for {request.url.path}")
        return response
    except NetworkError as e:
        logger.error(f"Network error caught in middleware: {e}")
        raise HTTPException(
            status_code=503,
            detail="Environment unavailable. Please check if the Superset instance is running."
        )
# [/DEF:log_requests:Function]

# Include API routes
app.include_router(auth.router)
app.include_router(admin.router)
app.include_router(plugins.router, prefix="/api/plugins", tags=["Plugins"])
app.include_router(tasks.router, prefix="/api/tasks", tags=["Tasks"])
app.include_router(settings.router, prefix="/api/settings", tags=["Settings"])
app.include_router(connections.router, prefix="/api/settings/connections", tags=["Connections"])
app.include_router(environments.router, prefix="/api/environments", tags=["Environments"])
app.include_router(mappings.router)
app.include_router(migration.router)
app.include_router(git.router)
app.include_router(llm.router)
app.include_router(storage.router, prefix="/api/storage", tags=["Storage"])

# [DEF:websocket_endpoint:Function]
# @PURPOSE: Provides a WebSocket endpoint for real-time log streaming of a task.
# @PRE: task_id must be a valid task ID.
# @POST: WebSocket connection is managed and logs are streamed until disconnect.
@app.websocket("/ws/logs/{task_id}")
async def websocket_endpoint(websocket: WebSocket, task_id: str):
    with belief_scope("websocket_endpoint", f"task_id={task_id}"):
        await websocket.accept()
    logger.info(f"WebSocket connection accepted for task {task_id}")
    task_manager = get_task_manager()
    queue = await task_manager.subscribe_logs(task_id)
    try:
        # Stream new logs
        logger.info(f"Starting log stream for task {task_id}")
        
        # Send initial logs first to build context
        initial_logs = task_manager.get_task_logs(task_id)
        for log_entry in initial_logs:
            log_dict = log_entry.dict()
            log_dict['timestamp'] = log_dict['timestamp'].isoformat()
            await websocket.send_json(log_dict)

        # Force a check for AWAITING_INPUT status immediately upon connection
        # This ensures that if the task is already waiting when the user connects, they get the prompt.
        task = task_manager.get_task(task_id)
        if task and task.status == "AWAITING_INPUT" and task.input_request:
             # Construct a synthetic log entry to trigger the frontend handler
             # This is a bit of a hack but avoids changing the websocket protocol significantly
             synthetic_log = {
                 "timestamp": task.logs[-1].timestamp.isoformat() if task.logs else "2024-01-01T00:00:00",
                 "level": "INFO",
                 "message": "Task paused for user input (Connection Re-established)",
                 "context": {"input_request": task.input_request}
             }
             await websocket.send_json(synthetic_log)

        while True:
            log_entry = await queue.get()
            log_dict = log_entry.dict()
            log_dict['timestamp'] = log_dict['timestamp'].isoformat()
            await websocket.send_json(log_dict)
            
            # If task is finished, we could potentially close the connection
            # but let's keep it open for a bit or until the client disconnects
            if "Task completed successfully" in log_entry.message or "Task failed" in log_entry.message:
                # Wait a bit to ensure client receives the last message
                await asyncio.sleep(2)
                # DO NOT BREAK here - allow client to keep connection open if they want to review logs
                # or until they disconnect. Breaking closes the socket immediately.
                # break

    except WebSocketDisconnect:
        logger.info(f"WebSocket connection disconnected for task {task_id}")
    except Exception as e:
        logger.error(f"WebSocket error for task {task_id}: {e}")
    finally:
        task_manager.unsubscribe_logs(task_id, queue)
# [/DEF:websocket_endpoint:Function]

# [DEF:StaticFiles:Mount]
# @SEMANTICS: static, frontend, spa
# @PURPOSE: Mounts the frontend build directory to serve static assets.
frontend_path = project_root / "frontend" / "build"
if frontend_path.exists():
    app.mount("/_app", StaticFiles(directory=str(frontend_path / "_app")), name="static")
    
    # Serve other static files from the root of build directory
    # [DEF:serve_spa:Function]
    # @PURPOSE: Serves frontend static files or index.html for SPA routing.
    # @PRE: file_path is requested by the client.
    # @POST: Returns the requested file or index.html as a fallback.
    @app.get("/{file_path:path}")
    async def serve_spa(file_path: str):
        with belief_scope("serve_spa", f"path={file_path}"):
            # Don't serve SPA for API routes that fell through
            if file_path.startswith("api/"):
                logger.info(f"[DEBUG] API route fell through to serve_spa: {file_path}")
                raise HTTPException(status_code=404, detail=f"API endpoint not found: {file_path}")
                
            full_path = frontend_path / file_path
            if full_path.is_file():
                return FileResponse(str(full_path))
            # Fallback to index.html for SPA routing
            return FileResponse(str(frontend_path / "index.html"))
    # [/DEF:serve_spa:Function]
else:
    # [DEF:read_root:Function]
    # @PURPOSE: A simple root endpoint to confirm that the API is running when frontend is missing.
    # @PRE: None.
    # @POST: Returns a JSON message indicating API status.
    @app.get("/")
    async def read_root():
        with belief_scope("read_root"):
            return {"message": "Superset Tools API is running (Frontend build not found)"}
    # [/DEF:read_root:Function]
# [/DEF:StaticFiles:Mount]
# [/DEF:AppModule:Module]
