# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
FastAPI application for the Satellite Environment.

This module creates an HTTP server that exposes the SatelliteEnvironment
over HTTP and WebSocket endpoints, compatible with EnvClient.

Endpoints:
    - POST /reset: Reset the environment
    - POST /step: Execute an action
    - GET /state: Get current environment state
    - GET /schema: Get action/observation schemas
    - WS /ws: WebSocket endpoint for persistent sessions

Usage:
    # Development (with auto-reload):
    uvicorn server.app:app --reload --host 0.0.0.0 --port 8000

    # Production:
    uvicorn server.app:app --host 0.0.0.0 --port 8000 --workers 4

    # Or run directly:
    python -m server.app
"""

try:
    from openenv.core.env_server.http_server import create_app
except Exception as e:  # pragma: no cover
    raise ImportError(
        "openenv-core[core] is required for the web interface. Install dependencies with "
        "'\n    uv pip install openenv-core[core] fastapi uvicorn\n'"
    ) from e

import importlib
from pathlib import Path

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse
from fastapi.openapi.utils import get_openapi
from fastapi.staticfiles import StaticFiles

# 1) Try absolute imports first (this handles the Docker layout where /app/env
#    is on PYTHONPATH and modules like `models` and `server` are top-level).
try:
    from models import SatelliteAction, SatelliteObservation
    from server.satellite_environment import SatelliteEnvironment
    from server.ui_demo import SatelliteUIDemo
except Exception:  # pragma: no cover - tolerate multiple packaging layouts
    # 2) Try package-relative imports (when running as `satellite.server.app`).
    try:
        from ..models import SatelliteAction, SatelliteObservation
        from .satellite_environment import SatelliteEnvironment
        from .ui_demo import SatelliteUIDemo
    except Exception:
        # 3) Try common installed package names (e.g., `satellite` or `env`).
        SatelliteAction = None
        SatelliteObservation = None
        SatelliteEnvironment = None
        SatelliteUIDemo = None
        for pkg in ("satellite", "env"):
            try:
                mod_models = importlib.import_module(f"{pkg}.models")
                mod_server_env = importlib.import_module(f"{pkg}.server.satellite_environment")
                mod_ui_demo = importlib.import_module(f"{pkg}.server.ui_demo")
                SatelliteAction = getattr(mod_models, "SatelliteAction")
                SatelliteObservation = getattr(mod_models, "SatelliteObservation")
                SatelliteEnvironment = getattr(mod_server_env, "SatelliteEnvironment")
                SatelliteUIDemo = getattr(mod_ui_demo, "SatelliteUIDemo")
                break
            except Exception:
                continue

        if SatelliteAction is None:
            # 4) Last resort: file-based fallback (load sibling files directly).
            try:
                import importlib.util
                from pathlib import Path

                here = Path(__file__).resolve().parent
                pkg_root = here.parent
                models_path = pkg_root / "models.py"
                server_env_path = here / "satellite_environment.py"
                ui_demo_path = here / "ui_demo.py"

                if models_path.exists() and server_env_path.exists() and ui_demo_path.exists():
                    spec_models = importlib.util.spec_from_file_location(
                        "satellite_models_fallback", str(models_path)
                    )
                    mod_models = importlib.util.module_from_spec(spec_models)  # type: ignore[arg-type]
                    spec_models.loader.exec_module(mod_models)  # type: ignore[attr-defined]

                    spec_server_env = importlib.util.spec_from_file_location(
                        "satellite_server_env_fallback", str(server_env_path)
                    )
                    mod_server_env = importlib.util.module_from_spec(spec_server_env)  # type: ignore[arg-type]
                    spec_server_env.loader.exec_module(mod_server_env)  # type: ignore[attr-defined]

                    spec_ui_demo = importlib.util.spec_from_file_location(
                        "satellite_ui_demo_fallback", str(ui_demo_path)
                    )
                    mod_ui_demo = importlib.util.module_from_spec(spec_ui_demo)  # type: ignore[arg-type]
                    spec_ui_demo.loader.exec_module(mod_ui_demo)  # type: ignore[attr-defined]

                    SatelliteAction = getattr(mod_models, "SatelliteAction")
                    SatelliteObservation = getattr(mod_models, "SatelliteObservation")
                    SatelliteEnvironment = getattr(mod_server_env, "SatelliteEnvironment")
                    SatelliteUIDemo = getattr(mod_ui_demo, "SatelliteUIDemo")
                else:
                    raise
            except Exception:
                raise


# Build the OpenEnv API app first, then expose its non-UI routes from our
# top-level FastAPI app so Hugging Face Spaces can use the custom frontend at `/`.
openenv_app = create_app(
    SatelliteEnvironment,
    SatelliteAction,
    SatelliteObservation,
    env_name="satellite",
    max_concurrent_envs=1,  # increase this number to allow more concurrent WebSocket sessions
)

app = FastAPI(
    title="Satellite Mission Tracking",
    version="0.1.0",
    description="Custom satellite mission frontend with mounted OpenEnv API routes.",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=[
        "http://localhost:4173",
        "http://127.0.0.1:4173",
        "http://localhost:5173",
        "http://127.0.0.1:5173",
    ],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

for route in openenv_app.router.routes:
    if getattr(route, "path", None) in {"/", "/docs", "/redoc", "/openapi.json"}:
        continue
    app.router.routes.append(route)

ui_demo = SatelliteUIDemo()
frontend_dist = Path(__file__).resolve().parent.parent / "frontend" / "dist"

if frontend_dist.exists():
    assets_dir = frontend_dist / "assets"
    if assets_dir.exists():
        app.mount("/assets", StaticFiles(directory=str(assets_dir)), name="assets")
        app.mount("/web/assets", StaticFiles(directory=str(assets_dir)), name="web-assets")


@app.get("/", include_in_schema=False)
def serve_root_index():
    index_path = frontend_dist / "index.html"
    if not index_path.exists():
        raise HTTPException(status_code=404, detail="Frontend build not found. Build /frontend first.")
    return FileResponse(index_path)


@app.get("/web", include_in_schema=False)
def serve_web_index():
    index_path = frontend_dist / "index.html"
    if not index_path.exists():
        raise HTTPException(status_code=404, detail="Frontend build not found. Build /frontend first.")
    return FileResponse(index_path)


@app.get("/web/{full_path:path}", include_in_schema=False)
def serve_web_app(full_path: str):
    candidate = frontend_dist / full_path
    index_path = frontend_dist / "index.html"
    if candidate.exists() and candidate.is_file():
        return FileResponse(candidate)
    if index_path.exists():
        return FileResponse(index_path)
    raise HTTPException(status_code=404, detail="Frontend build not found. Build /frontend first.")


@app.get("/{full_path:path}", include_in_schema=False)
def serve_frontend_app(full_path: str):
    if full_path.startswith(("api/", "docs", "openapi.json", "redoc", "ws", "reset", "step", "state", "schema", "health")):
        raise HTTPException(status_code=404, detail="Not found")
    candidate = frontend_dist / full_path
    index_path = frontend_dist / "index.html"
    if candidate.exists() and candidate.is_file():
        return FileResponse(candidate)
    if index_path.exists():
        return FileResponse(index_path)
    raise HTTPException(status_code=404, detail="Frontend build not found. Build /frontend first.")


@app.get("/api/ui/demo")
def ui_demo_snapshot():
    return ui_demo.snapshot()


@app.post("/api/ui/demo/reset")
def ui_demo_reset(task_name: str = "medium"):
    if task_name not in {"easy", "medium", "hard"}:
        raise HTTPException(status_code=400, detail="task_name must be easy, medium, or hard")
    return ui_demo.reset(task_name)


@app.post("/api/ui/demo/step")
def ui_demo_step():
    return ui_demo.step()


def _with_examples() -> None:
    """Inject endpoint examples into Swagger / OpenAPI docs."""
    step_action = {
        "action": {
            "satellite_actions": {
                "0": "capture",
                "1": "maintain",
                "2": "idle",
            }
        },
        "timeout_s": 30,
    }
    reset_request = {"task_name": "easy", "timeout_s": 30}
    observation_example = {
        "satellites": [
            {
                "id": 0,
                "position": [7000.0, 0.0, 0.0],
                "battery": 82.5,
                "storage": 35.0,
                "last_action": "capture",
            },
            {
                "id": 1,
                "position": [0.0, 7000.0, 0.0],
                "battery": 91.0,
                "storage": 10.0,
                "last_action": "maintain",
            },
        ],
        "time_step": 1,
        "ground_stations": [[28.6139, 77.2090], [34.0522, -118.2437]],
        "weather_conditions": {"region_0": 0.2, "region_1": 0.6},
        "pending_tasks": [
            {
                "id": "img_001",
                "type": "capture",
                "target": "region_0",
                "priority": 0.9,
            }
        ],
        "total_reward": 4.5,
        "done": False,
        "reward": 4.5,
        "metadata": {
            "task_name": "easy",
            "reward_components": {"capture": 5.0, "idle_penalty": -0.5},
        },
    }
    step_response = {
        "observation": observation_example,
        "reward": 4.5,
        "done": False,
        "info": {
            "task_name": "easy",
            "reward_components": {"capture": 5.0, "idle_penalty": -0.5},
        },
    }
    state_response = {
        "episode_id": "demo-episode-123",
        "step_count": 1,
    }
    schema_response = {
        "action_schema": {
            "title": "SatelliteAction",
            "type": "object",
            "properties": {
                "satellite_actions": {
                    "type": "object",
                    "additionalProperties": {
                        "enum": ["capture", "downlink", "maintain", "idle"]
                    },
                }
            },
        },
        "observation_schema": {
            "title": "SatelliteObservation",
            "type": "object",
            "properties": {
                "satellites": {"type": "array"},
                "time_step": {"type": "integer"},
                "ground_stations": {"type": "array"},
                "weather_conditions": {"type": "object"},
                "pending_tasks": {"type": "array"},
                "total_reward": {"type": "number"},
                "done": {"type": "boolean"},
                "reward": {"type": "number"},
                "metadata": {"type": "object"},
            },
        },
    }

    def custom_openapi():
        if app.openapi_schema:
            return app.openapi_schema

        schema = get_openapi(
            title=app.title or "Satellite OpenEnv API",
            version=app.version or "0.1.0",
            description=app.description or "HTTP API for the Satellite OpenEnv environment.",
            routes=app.routes,
        )

        paths = schema.setdefault("paths", {})

        reset_post = paths.get("/reset", {}).get("post")
        if reset_post:
            content = reset_post.setdefault("requestBody", {}).setdefault("content", {})
            content.setdefault("application/json", {})["example"] = reset_request
            responses = reset_post.setdefault("responses", {})
            ok = responses.setdefault("200", {}).setdefault("content", {})
            ok.setdefault("application/json", {})["example"] = observation_example

        step_post = paths.get("/step", {}).get("post")
        if step_post:
            content = step_post.setdefault("requestBody", {}).setdefault("content", {})
            content.setdefault("application/json", {})["example"] = step_action
            responses = step_post.setdefault("responses", {})
            ok = responses.setdefault("200", {}).setdefault("content", {})
            ok.setdefault("application/json", {})["example"] = step_response

        state_get = paths.get("/state", {}).get("get")
        if state_get:
            responses = state_get.setdefault("responses", {})
            ok = responses.setdefault("200", {}).setdefault("content", {})
            ok.setdefault("application/json", {})["example"] = state_response

        schema_get = paths.get("/schema", {}).get("get")
        if schema_get:
            responses = schema_get.setdefault("responses", {})
            ok = responses.setdefault("200", {}).setdefault("content", {})
            ok.setdefault("application/json", {})["example"] = schema_response

        app.openapi_schema = schema
        return app.openapi_schema

    app.openapi = custom_openapi


_with_examples()


def run_server(host: str = "0.0.0.0", port: int = 8000):
    """
    Entry point for direct execution via uv run or python -m.

    This function enables running the server without Docker:
        uv run --project . server
        uv run --project . server --port 8001
        python -m satellite.server.app

    Args:
        host: Host address to bind to (default: "0.0.0.0")
        port: Port number to listen on (default: 8000)

    For production deployments, consider using uvicorn directly with
    multiple workers:
        uvicorn satellite.server.app:app --workers 4
    """
    import uvicorn

    uvicorn.run(app, host=host, port=port)


def main():
    """CLI-compatible main entry point for OpenEnv validation and local runs."""
    import argparse

    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8000)
    args = parser.parse_args()
    run_server(host=args.host, port=args.port)


if __name__ == "__main__":
    main()