"""
NeuroFlow Backend - Main FastAPI Application

A local-first, node-based ML orchestrator API.
"""

from __future__ import annotations

import asyncio
import json
import os
from contextlib import asynccontextmanager
from datetime import datetime
from pathlib import Path
from typing import Any
from uuid import uuid4

from fastapi import (
    FastAPI,
    HTTPException,
    WebSocket,
    WebSocketDisconnect,
    UploadFile,
    File,
    BackgroundTasks,
)
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse

from .schemas import (
    APIResponse,
    ExecutionProgress,
    ExecutionRequest,
    ExecutionResult,
    ExecutionStatus,
    HealthResponse,
    Workflow,
    WorkflowCreate,
)
from .engine import ExecutionEngine, DAGValidationError
from .storage import StorageManager

# Import executors
from .executors.data_loaders import csv_loader_executor, json_loader_executor
from .executors.preprocessors import (
    data_scaler_executor,
    train_test_split_executor,
    missing_value_handler_executor,
    encoder_executor,
)
from .executors.trainers import trainer_executor
from .executors.evaluators import evaluator_executor, prediction_executor


# ============== Application Setup ==============

# Global instances
storage: StorageManager = None  # type: ignore
engine: ExecutionEngine = None  # type: ignore

# WebSocket connections for progress updates
active_connections: dict[str, list[WebSocket]] = {}


@asynccontextmanager
async def lifespan(app: FastAPI):
    """Application lifespan manager."""
    global storage, engine
    
    # Startup
    storage = StorageManager(base_path="./storage")
    engine = ExecutionEngine(storage=storage)
    
    # Register all executors
    engine.register_executor("csv_loader", csv_loader_executor)
    engine.register_executor("json_loader", json_loader_executor)
    engine.register_executor("data_scaler", data_scaler_executor)
    engine.register_executor("train_test_split", train_test_split_executor)
    engine.register_executor("missing_value_handler", missing_value_handler_executor)
    engine.register_executor("encoder", encoder_executor)
    engine.register_executor("trainer", trainer_executor)
    engine.register_executor("evaluator", evaluator_executor)
    engine.register_executor("prediction", prediction_executor)
    
    print("🚀 NeuroFlow Backend started")
    print(f"📁 Storage path: {storage.base_path.absolute()}")
    
    yield
    
    # Shutdown
    print("👋 NeuroFlow Backend shutting down")


# Create FastAPI app
app = FastAPI(
    title="NeuroFlow",
    description="A local-first, node-based ML orchestrator",
    version="0.1.0",
    lifespan=lifespan,
)

# CORS middleware for React frontend (development)
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:3000", "http://127.0.0.1:3000", "*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============== Static File Serving ==============

def get_static_dir() -> Path:
    """Get the path to static files directory."""
    # Check environment variable first (set by CLI)
    env_path = os.environ.get("NEUROFLOW_STATIC_DIR")
    if env_path:
        return Path(env_path)
    
    # Check common locations
    locations = [
        Path(__file__).parent.parent / "static",
        Path(__file__).parent.parent.parent / "static",
    ]
    
    for loc in locations:
        if loc.exists() and (loc / "index.html").exists():
            return loc
    
    return locations[0]


# Mount static files if they exist
_static_dir = get_static_dir()
if _static_dir.exists():
    # Serve _next folder for Next.js assets
    _next_dir = _static_dir / "_next"
    if _next_dir.exists():
        app.mount("/_next", StaticFiles(directory=str(_next_dir)), name="next_static")
    
    # Serve other static assets
    _images_dir = _static_dir / "images"
    if _images_dir.exists():
        app.mount("/images", StaticFiles(directory=str(_images_dir)), name="images")


# ============== Health Check ==============

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """Check API health status."""
    return HealthResponse(
        status="healthy",
        version="0.1.0",
    )


@app.get("/api", response_model=APIResponse)
@app.get("/api/", response_model=APIResponse)
async def api_root():
    """API root endpoint."""
    return APIResponse(
        success=True,
        message="Welcome to NeuroFlow API",
        data={"version": "0.1.0"},
    )


@app.get("/", response_class=HTMLResponse)
async def root():
    """Root endpoint - serve frontend or API info."""
    static_dir = get_static_dir()
    index_path = static_dir / "index.html"
    
    if index_path.exists():
        return FileResponse(str(index_path), media_type="text/html")
    
    # Fallback when frontend not built
    return HTMLResponse("""
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>NeuroFlow</title>
    <style>
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            display: flex;
            justify-content: center;
            align-items: center;
            min-height: 100vh;
            margin: 0;
            background: linear-gradient(135deg, #1a1a2e 0%, #16213e 100%);
            color: #fff;
        }
        .container { text-align: center; max-width: 600px; padding: 40px; }
        h1 { font-size: 3rem; margin-bottom: 1rem; }
        p { color: #a0a0a0; margin: 0.5rem 0; }
        code { background: #2d2d4a; padding: 4px 12px; border-radius: 4px; }
        .api-link {
            display: inline-block;
            margin-top: 2rem;
            padding: 12px 24px;
            background: #6366f1;
            color: white;
            text-decoration: none;
            border-radius: 8px;
            font-weight: 600;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>🧠 NeuroFlow</h1>
        <p>The backend is running, but the frontend static files were not found.</p>
        <p>To build the frontend:</p>
        <p><code>cd Frontend && pnpm build</code></p>
        <a href="/docs" class="api-link">View API Documentation</a>
    </div>
</body>
</html>
    """, status_code=200)


# ============== Workflow Endpoints ==============

@app.post("/api/workflows/validate", response_model=APIResponse)
async def validate_workflow(workflow: Workflow):
    """
    Validate a workflow DAG structure.
    
    Checks for cycles and invalid edge references.
    """
    try:
        is_valid, error = engine.validate_dag(workflow)
        
        if is_valid:
            execution_order = engine.topological_sort(workflow)
            return APIResponse(
                success=True,
                message="Workflow is valid",
                data={
                    "valid": True,
                    "execution_order": execution_order,
                    "node_count": len(workflow.nodes),
                    "edge_count": len(workflow.edges),
                },
            )
        else:
            return APIResponse(
                success=False,
                message=error,
                data={"valid": False},
            )
    except DAGValidationError as e:
        return APIResponse(
            success=False,
            message=str(e),
            data={"valid": False},
        )


@app.post("/api/workflows/execute", response_model=APIResponse)
async def execute_workflow(
    request: ExecutionRequest,
    background_tasks: BackgroundTasks,
):
    """
    Execute a workflow asynchronously.
    
    Returns immediately with a run_id that can be used
    to track progress via WebSocket or polling.
    """
    run_id = request.run_id or str(uuid4())
    workflow = request.workflow
    
    # Validate first
    is_valid, error = engine.validate_dag(workflow)
    if not is_valid:
        raise HTTPException(status_code=400, detail=error)
    
    # Create initial result
    initial_result = ExecutionResult(
        run_id=run_id,
        workflow_id=workflow.id,
        status=ExecutionStatus.PENDING,
        started_at=datetime.utcnow(),
    )
    
    storage.save_run_result(run_id, initial_result)
    
    # Define progress callback for WebSocket
    async def progress_callback(progress: ExecutionProgress):
        await broadcast_progress(run_id, progress)
    
    # Run execution in background
    async def run_execution():
        try:
            result = await engine.execute_workflow(
                workflow=workflow,
                run_id=run_id,
                progress_callback=lambda p: asyncio.create_task(
                    broadcast_progress(run_id, p)
                ),
            )
            # Broadcast completion
            await broadcast_progress(
                run_id,
                ExecutionProgress(
                    run_id=run_id,
                    workflow_id=workflow.id,
                    current_node_id="",
                    current_node_index=len(workflow.nodes),
                    total_nodes=len(workflow.nodes),
                    status=result.status,
                    message="Workflow execution completed",
                ),
            )
        except Exception as e:
            await broadcast_progress(
                run_id,
                ExecutionProgress(
                    run_id=run_id,
                    workflow_id=workflow.id,
                    current_node_id="",
                    current_node_index=0,
                    total_nodes=len(workflow.nodes),
                    status=ExecutionStatus.FAILED,
                    message=f"Execution failed: {e}",
                ),
            )
    
    # Add to background tasks
    background_tasks.add_task(asyncio.create_task, run_execution())
    
    return APIResponse(
        success=True,
        message="Workflow execution started",
        data={
            "run_id": run_id,
            "workflow_id": workflow.id,
            "status": "pending",
        },
    )


@app.post("/api/workflows/execute/sync", response_model=ExecutionResult)
async def execute_workflow_sync(request: ExecutionRequest):
    """
    Execute a workflow synchronously.
    
    Waits for completion and returns the full result.
    Use for small workflows or testing.
    """
    run_id = request.run_id or str(uuid4())
    workflow = request.workflow
    
    # Validate
    is_valid, error = engine.validate_dag(workflow)
    if not is_valid:
        raise HTTPException(status_code=400, detail=error)
    
    # Execute
    result = await engine.execute_workflow(
        workflow=workflow,
        run_id=run_id,
    )
    
    return result


# ============== Run Management Endpoints ==============

@app.get("/api/runs", response_model=APIResponse)
async def list_runs():
    """List all execution runs."""
    runs = storage.list_runs()
    return APIResponse(
        success=True,
        message=f"Found {len(runs)} runs",
        data=runs,
    )


@app.get("/api/runs/{run_id}", response_model=APIResponse)
async def get_run(run_id: str):
    """Get details of a specific run."""
    try:
        result = storage.load_run_result(run_id)
        return APIResponse(
            success=True,
            message="Run found",
            data=result,
        )
    except FileNotFoundError:
        raise HTTPException(status_code=404, detail=f"Run not found: {run_id}")


@app.get("/api/runs/{run_id}/logs", response_model=APIResponse)
async def get_run_logs(run_id: str):
    """Get execution logs for a run."""
    logs = storage.get_run_logs(run_id)
    return APIResponse(
        success=True,
        message=f"Found {len(logs)} log entries",
        data=logs,
    )


@app.delete("/api/runs/{run_id}", response_model=APIResponse)
async def delete_run(run_id: str):
    """Delete a run and all its artifacts."""
    deleted = storage.delete_run(run_id)
    if deleted:
        return APIResponse(
            success=True,
            message=f"Run {run_id} deleted",
        )
    else:
        raise HTTPException(status_code=404, detail=f"Run not found: {run_id}")


@app.post("/api/runs/{run_id}/cancel", response_model=APIResponse)
async def cancel_run(run_id: str):
    """Cancel a running workflow execution."""
    cancelled = engine.cancel_execution(run_id)
    if cancelled:
        return APIResponse(
            success=True,
            message=f"Run {run_id} cancelled",
        )
    else:
        return APIResponse(
            success=False,
            message=f"Run {run_id} not found or already completed",
        )


# ============== Model Endpoints ==============

@app.get("/api/runs/{run_id}/models", response_model=APIResponse)
async def list_run_models(run_id: str):
    """List all models saved during a run."""
    models = storage.list_models(run_id)
    return APIResponse(
        success=True,
        message=f"Found {len(models)} models",
        data=models,
    )


@app.get("/api/runs/{run_id}/models/{model_name}/download")
async def download_model(run_id: str, model_name: str):
    """Download a trained model file."""
    try:
        model_path = storage._get_models_path(run_id) / f"{model_name}.joblib"
        if not model_path.exists():
            raise HTTPException(status_code=404, detail="Model not found")
        
        return FileResponse(
            path=str(model_path),
            filename=f"{model_name}.joblib",
            media_type="application/octet-stream",
        )
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# ============== Data File Endpoints ==============

@app.get("/api/data", response_model=APIResponse)
async def list_data_files():
    """List all available data files."""
    files = storage.list_data_files()
    return APIResponse(
        success=True,
        message=f"Found {len(files)} data files",
        data=files,
    )


@app.post("/api/data/upload", response_model=APIResponse)
async def upload_data_file(file: UploadFile = File(...)):
    """Upload a data file."""
    if not file.filename:
        raise HTTPException(status_code=400, detail="No filename provided")
    
    content = await file.read()
    file_path = storage.save_uploaded_file(file.filename, content)
    
    return APIResponse(
        success=True,
        message=f"File uploaded: {file.filename}",
        data={"path": str(file_path)},
    )


# ============== WebSocket for Progress Updates ==============

async def broadcast_progress(run_id: str, progress: ExecutionProgress):
    """Broadcast progress update to all connected clients for a run."""
    if run_id not in active_connections:
        return
    
    disconnected = []
    
    for websocket in active_connections[run_id]:
        try:
            await websocket.send_json(progress.model_dump(mode="json"))
        except Exception:
            disconnected.append(websocket)
    
    # Remove disconnected clients
    for ws in disconnected:
        active_connections[run_id].remove(ws)


@app.websocket("/ws/runs/{run_id}")
async def websocket_run_progress(websocket: WebSocket, run_id: str):
    """
    WebSocket endpoint for real-time execution progress.
    
    Connect to receive progress updates for a specific run.
    """
    await websocket.accept()
    
    # Add to active connections
    if run_id not in active_connections:
        active_connections[run_id] = []
    active_connections[run_id].append(websocket)
    
    try:
        # Send initial status
        try:
            result = storage.load_run_result(run_id)
            await websocket.send_json({
                "type": "initial",
                "data": result,
            })
        except FileNotFoundError:
            await websocket.send_json({
                "type": "initial",
                "data": {"status": "not_found"},
            })
        
        # Keep connection open
        while True:
            try:
                # Wait for any client messages (like ping)
                data = await websocket.receive_text()
                
                if data == "ping":
                    await websocket.send_text("pong")
                elif data == "status":
                    try:
                        result = storage.load_run_result(run_id)
                        await websocket.send_json({
                            "type": "status",
                            "data": result,
                        })
                    except FileNotFoundError:
                        pass
                        
            except WebSocketDisconnect:
                break
                
    finally:
        # Remove from active connections
        if run_id in active_connections:
            if websocket in active_connections[run_id]:
                active_connections[run_id].remove(websocket)
            if not active_connections[run_id]:
                del active_connections[run_id]


# ============== Static File Serving ==============

# Mount static files for React frontend (if exists)
import os
if os.path.exists("./static"):
    app.mount("/static", StaticFiles(directory="static"), name="static")
    
    @app.get("/app/{full_path:path}")
    async def serve_spa(full_path: str):
        """Serve the React SPA for any /app route."""
        return FileResponse("./static/index.html")


# ============== Node Types Info ==============

@app.get("/api/node-types", response_model=APIResponse)
async def get_node_types():
    """Get information about all available node types."""
    node_types = {
        "data_loaders": {
            "csv_loader": {
                "label": "CSV Loader",
                "description": "Load data from a CSV file",
                "params": {
                    "file_path": {"type": "string", "required": True},
                    "delimiter": {"type": "string", "default": ","},
                    "has_header": {"type": "boolean", "default": True},
                    "encoding": {"type": "string", "default": "utf-8"},
                },
                "inputs": [],
                "outputs": ["dataframe"],
            },
            "json_loader": {
                "label": "JSON Loader",
                "description": "Load data from a JSON file",
                "params": {
                    "file_path": {"type": "string", "required": True},
                    "orient": {"type": "string", "default": "records"},
                },
                "inputs": [],
                "outputs": ["dataframe"],
            },
        },
        "preprocessors": {
            "data_scaler": {
                "label": "Data Scaler",
                "description": "Scale numerical features",
                "params": {
                    "scaler_type": {
                        "type": "enum",
                        "options": ["standard", "minmax", "robust", "normalizer"],
                        "default": "standard",
                    },
                    "columns": {"type": "array", "default": []},
                },
                "inputs": ["dataframe"],
                "outputs": ["dataframe"],
            },
            "train_test_split": {
                "label": "Train/Test Split",
                "description": "Split data into training and test sets",
                "params": {
                    "test_size": {"type": "number", "default": 0.2},
                    "random_state": {"type": "integer", "default": 42},
                    "target_column": {"type": "string"},
                    "shuffle": {"type": "boolean", "default": True},
                },
                "inputs": ["dataframe"],
                "outputs": ["train_data", "test_data"],
            },
            "missing_value_handler": {
                "label": "Missing Value Handler",
                "description": "Handle missing values in data",
                "params": {
                    "strategy": {
                        "type": "enum",
                        "options": ["drop", "mean", "median", "mode", "constant"],
                        "default": "drop",
                    },
                    "fill_value": {"type": "any", "default": 0},
                },
                "inputs": ["dataframe"],
                "outputs": ["dataframe"],
            },
            "encoder": {
                "label": "Encoder",
                "description": "Encode categorical variables",
                "params": {
                    "encoding_type": {
                        "type": "enum",
                        "options": ["label", "onehot"],
                        "default": "label",
                    },
                    "columns": {"type": "array", "default": []},
                },
                "inputs": ["dataframe"],
                "outputs": ["dataframe"],
            },
        },
        "trainers": {
            "trainer": {
                "label": "Model Trainer",
                "description": "Train an ML model",
                "params": {
                    "model_type": {
                        "type": "enum",
                        "options": [
                            "logistic_regression",
                            "random_forest_classifier",
                            "svm_classifier",
                            "knn_classifier",
                            "linear_regression",
                            "random_forest_regressor",
                            "gradient_boosting_classifier",
                            "gradient_boosting_regressor",
                        ],
                        "default": "random_forest_classifier",
                    },
                    "target_column": {"type": "string"},
                    "hyperparameters": {"type": "object", "default": {}},
                },
                "inputs": ["dataframe"],
                "outputs": ["model"],
            },
        },
        "evaluators": {
            "evaluator": {
                "label": "Model Evaluator",
                "description": "Evaluate model performance",
                "params": {
                    "metrics": {
                        "type": "array",
                        "default": ["accuracy", "f1"],
                        "options": ["accuracy", "precision", "recall", "f1", "mse", "rmse", "mae", "r2"],
                    },
                    "cross_validation": {"type": "boolean", "default": False},
                    "cv_folds": {"type": "integer", "default": 5},
                },
                "inputs": ["model"],
                "outputs": ["metrics"],
            },
            "prediction": {
                "label": "Prediction",
                "description": "Make predictions with a trained model",
                "params": {},
                "inputs": ["model", "dataframe"],
                "outputs": ["predictions"],
            },
        },
    }
    
    return APIResponse(
        success=True,
        message="Node types retrieved",
        data=node_types,
    )


# ============== Catch-all route for SPA ==============

@app.get("/{full_path:path}")
async def serve_spa(full_path: str):
    """Serve static files or fall back to index.html for client-side routing."""
    static_dir = get_static_dir()
    
    # Don't catch API routes
    if full_path.startswith("api/") or full_path.startswith("ws/"):
        raise HTTPException(status_code=404, detail="Not found")
    
    # Try to serve the exact file
    file_path = static_dir / full_path
    if file_path.exists() and file_path.is_file():
        return FileResponse(str(file_path))
    
    # Fall back to index.html for SPA routing
    index_path = static_dir / "index.html"
    if index_path.exists():
        return FileResponse(str(index_path), media_type="text/html")
    
    raise HTTPException(status_code=404, detail="Not found")


if __name__ == "__main__":
    import uvicorn
    uvicorn.run(
        "app.main:app",
        host="0.0.0.0",
        port=8000,
        reload=True,
    )
