# 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 Constraint Env Environment.

This module creates an HTTP server that exposes the ConstraintEnvironment
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
"""

import os


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

try:
    from ..models import ConstraintAction, ConstraintObservation
    from .constraint_env_environment import ConstraintEnvironment
except ImportError:
    from constraint_env.models import ConstraintAction, ConstraintObservation
    from constraint_env.server.constraint_env_environment import ConstraintEnvironment

# Load the dataset so the environment can be initialised without crashing.
try:
    from dataset_example import dataset as _DATASET
except ImportError:
    from constraint_env.dataset_example import dataset as _DATASET  # type: ignore


def _make_env():
    """Factory that passes the pre-loaded dataset into ConstraintEnvironment."""
    return ConstraintEnvironment(dataset=_DATASET)


try:
    from .gradio_ui import build_constraint_gradio_ui as _gradio_builder
except ImportError:
    try:
        from constraint_env.server.gradio_ui import build_constraint_gradio_ui as _gradio_builder
    except ImportError:
        _gradio_builder = None

# Create the app – pass the factory so create_app calls _make_env() per session.
import gradio as gr
_orig_tabbed = gr.TabbedInterface

def _swapped_tabbed(interface_list, tab_names, **kwargs):
    if len(interface_list) == 2 and "Custom" in tab_names:
        idx_custom = tab_names.index("Custom")
        idx_play = tab_names.index("Playground")
        return _orig_tabbed(
            [interface_list[idx_custom], interface_list[idx_play]],
            ["Constraint Compiler", "Base Playground"],
            **kwargs
        )
    return _orig_tabbed(interface_list, tab_names, **kwargs)

gr.TabbedInterface = _swapped_tabbed

import openenv.core.env_server.web_interface as wi

wi._load_readme_from_filesystem = lambda *args: """
An interactive compiler environment where AI agents (or humans!) translate natural language scheduling rules into rigorous JSON Abstract Syntax Trees (ASTs). 

### Objective
Read the natural language prompt, draft the corresponding JSON AST using the environment's strict DSL schema, and compile it. If you make a mistake, the environment will give you deterministic compiler feedback so you can debug your logic!

### How to Interact
1. **Reset** → Click `Reset` to load a fresh task and read the Prompt.
2. **Draft** → Write your JSON AST in the input box.
3. **Step** → Click `Step` to submit your AST to the OpenEnv compiler.
4. **Debug** → Check the Observation output. Read the feedback, fix your JSON, and click `Step` again!

### Rules of the Compiler
* **Valid JSON Only:** The compiler will immediately reject malformed JSON.
* **Declare Variables:** Every variable used in `where` or `assert` must be declared in the `forall` array.
* **Step Penalties:** Every time you submit an incorrect AST, you lose `-0.05` points from your maximum possible reward. Solve it in as few steps as possible!

###  Difficulty Levels
* **🟢 Easy** → Direct equality matches and simple assertions (Target: 0.80+).
* **🟡 Medium** → Aggregate objectives requiring `sum` operations and nested loops (Target: 0.65+).
* **🔴 Hard** → Deep conditional traversals with compound boolean logic and `NOT` filters (Target: 0.45+).

### Example Step JSON Input
```json
{
  "type": "hard",
  "name": "cs_department_meeting",
  "forall": [
    {"b": "branches"},
    {"sub": {"subjects": "b"}},
    {"d": "days"},
    {"s": "slots"}
  ],
  "where": {
    "operator": "AND",
    "left": {"operator": "==", "left": {"name": "b"}, "right": "CS"},
    "right": {"operator": "==", "left": "d", "right": 2}
  },
  "assert": {
    "operator": "==",
    "left": {
      "target": "schedule",
      "args": [{"name": "b"}, {"name": "sub"}, "d", "s"]
    },
    "right": 0
  }
}
```"""

app = create_app(
    _make_env,
    ConstraintAction,
    ConstraintObservation,
    env_name="Timetable Constraint Environment (NL-to-AST)",
    max_concurrent_envs=1,
    gradio_builder=_gradio_builder,
)


# ---------------------------------------------------------------------------
# PWA manifest – browsers request this at root level for the web UI
# ---------------------------------------------------------------------------

import os
from pathlib import Path
from fastapi.staticfiles import StaticFiles
from fastapi.responses import JSONResponse

_ASSETS_DIR = Path(__file__).parent.parent / "assets"
if _ASSETS_DIR.exists():
    app.mount("/assets", StaticFiles(directory=str(_ASSETS_DIR)), name="assets")


@app.get("/manifest.json", include_in_schema=False)
async def web_manifest():
    return JSONResponse(
        content={
            "name": "Timetable Constraint Environment (NL-to-AST)",
            "short_name": "NL-to-AST",
            "description": "RL training environment: natural-language → constraint AST",
            "start_url": "/web/",
            "display": "standalone",
            "background_color": "#1e1e2e",
            "theme_color": "#7c3aed",
            "icons": [
                {
                    "src": "https://huggingface.co/front/assets/huggingface_logo-noborder.svg",
                    "sizes": "any",
                    "type": "image/svg+xml",
                }
            ],
        },
        headers={"Cache-Control": "public, max-age=3600"},
    )


def main():
    """
    Entry point for direct execution via uv run or python -m.

    Callable with no arguments (required for openenv validate).
    Respects $PORT env var so it works on HF Spaces (7860) and locally.

        uv run server/app.py
        python -m constraint_env.server.app
    """
    import argparse
    import uvicorn

    parser = argparse.ArgumentParser(description="Constraint Env FastAPI Server")
    parser.add_argument("--host", type=str, default="0.0.0.0")
    parser.add_argument("--port", type=int, default=int(os.environ.get("PORT", 7860)))
    args = parser.parse_args()

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


if __name__ == "__main__":
    main()