#!/usr/bin/env python3
"""Public Optuna Dashboard: live DB with hard read-only enforcement.

Two layers (either alone is not enough if the wrong process is running):

1. HTTP: reject POST/PUT/DELETE/PATCH with 403.
2. Storage: wrap RDBStorage so delete_study / tell / set_* raise and never hit SQLite.

Plain SQLite ``?mode=ro`` is ignored by SQLAlchemy/Optuna — do not rely on it.
"""
from __future__ import annotations

import argparse
import sys
from socketserver import ThreadingMixIn
from typing import Any, Container, Sequence
from wsgiref.simple_server import WSGIServer, make_server

from optuna.storages import BaseStorage, RDBStorage
from optuna.study import StudyDirection
from optuna.trial import FrozenTrial, TrialState
from optuna_dashboard import wsgi

_ALLOWED_METHODS = frozenset({"GET", "HEAD", "OPTIONS"})

# BaseStorage / RDB methods that mutate data (block all of these).
_WRITE_METHODS = frozenset(
    {
        "create_new_study",
        "create_new_trial",
        "delete_study",
        "set_study_user_attr",
        "set_study_system_attr",
        "set_trial_param",
        "set_trial_state_values",
        "set_trial_intermediate_value",
        "set_trial_user_attr",
        "set_trial_system_attr",
        "remove_session",
    }
)


class ReadOnlyError(PermissionError):
    """Raised when a write is attempted on the public dashboard."""


class ReadOnlyStorage(BaseStorage):
    """Delegates reads to an inner storage; blocks every write API."""

    def __init__(self, inner: BaseStorage) -> None:
        self._inner = inner

    def _block(self, name: str) -> None:
        raise ReadOnlyError(
            f"read-only public dashboard: {name}() is disabled "
            "(cannot create/delete/modify studies or trials)"
        )

    # --- writes ---
    def create_new_study(
        self,
        directions: Sequence[StudyDirection],
        study_name: str | None = None,
    ) -> int:
        self._block("create_new_study")

    def delete_study(self, study_id: int) -> None:
        self._block("delete_study")

    def set_study_user_attr(self, study_id: int, key: str, value: Any) -> None:
        self._block("set_study_user_attr")

    def set_study_system_attr(self, study_id: int, key: str, value: Any) -> None:
        self._block("set_study_system_attr")

    def create_new_trial(self, study_id: int, template_trial: FrozenTrial | None = None) -> int:
        self._block("create_new_trial")

    def set_trial_param(
        self,
        trial_id: int,
        param_name: str,
        param_value_internal: float,
        distribution: Any,
    ) -> None:
        self._block("set_trial_param")

    def set_trial_state_values(
        self,
        trial_id: int,
        state: TrialState,
        values: Sequence[float] | None = None,
    ) -> bool:
        self._block("set_trial_state_values")

    def set_trial_intermediate_value(
        self, trial_id: int, step: int, intermediate_value: float
    ) -> None:
        self._block("set_trial_intermediate_value")

    def set_trial_user_attr(self, trial_id: int, key: str, value: Any) -> None:
        self._block("set_trial_user_attr")

    def set_trial_system_attr(self, trial_id: int, key: str, value: Any) -> None:
        self._block("set_trial_system_attr")

    # --- reads (delegate) ---
    def get_study_id_from_name(self, study_name: str) -> int:
        return self._inner.get_study_id_from_name(study_name)

    def get_study_name_from_id(self, study_id: int) -> str:
        return self._inner.get_study_name_from_id(study_id)

    def get_study_directions(self, study_id: int) -> list[StudyDirection]:
        return list(self._inner.get_study_directions(study_id))

    def get_study_user_attrs(self, study_id: int) -> dict[str, Any]:
        return self._inner.get_study_user_attrs(study_id)

    def get_study_system_attrs(self, study_id: int) -> dict[str, Any]:
        return self._inner.get_study_system_attrs(study_id)

    def get_all_studies(self) -> list[Any]:
        return self._inner.get_all_studies()

    def get_trial_id_from_study_id_trial_number(self, study_id: int, trial_number: int) -> int:
        return self._inner.get_trial_id_from_study_id_trial_number(study_id, trial_number)

    def get_trial_number_from_id(self, trial_id: int) -> int:
        return self._inner.get_trial_number_from_id(trial_id)

    def get_trial(self, trial_id: int) -> FrozenTrial:
        return self._inner.get_trial(trial_id)

    def get_all_trials(
        self,
        study_id: int,
        deepcopy: bool = True,
        states: Container[TrialState] | None = None,
    ) -> list[FrozenTrial]:
        return self._inner.get_all_trials(study_id, deepcopy=deepcopy, states=states)

    def get_n_trials(
        self, study_id: int, state: tuple[TrialState, ...] | TrialState | None = None
    ) -> int:
        return self._inner.get_n_trials(study_id, state=state)

    def get_best_trial(self, study_id: int) -> FrozenTrial:
        return self._inner.get_best_trial(study_id)

    def get_trial_param(self, trial_id: int, param_name: str) -> float:
        return self._inner.get_trial_param(trial_id, param_name)

    def get_trial_params(self, trial_id: int) -> dict[str, Any]:
        # optional on some versions
        if hasattr(self._inner, "get_trial_params"):
            return self._inner.get_trial_params(trial_id)  # type: ignore[attr-defined]
        return self.get_trial(trial_id).params

    def get_trial_user_attrs(self, trial_id: int) -> dict[str, Any]:
        return self._inner.get_trial_user_attrs(trial_id)

    def get_trial_system_attrs(self, trial_id: int) -> dict[str, Any]:
        return self._inner.get_trial_system_attrs(trial_id)

    def check_trial_is_updatable(self, trial_id: int, trial_state: TrialState) -> None:
        # Used before writes; fail closed.
        self._block("check_trial_is_updatable")

    def remove_session(self) -> None:
        # Session cleanup can be a no-op for read-only RDB use.
        if hasattr(self._inner, "remove_session"):
            try:
                self._inner.remove_session()  # type: ignore[attr-defined]
            except Exception:
                pass

    def __getattr__(self, name: str) -> Any:
        # Fail closed for any other mutators Optuna may add.
        if name in _WRITE_METHODS or name.startswith("set_") or name.startswith("delete_"):
            def _blocked(*_a: Any, **_k: Any) -> Any:
                self._block(name)

            return _blocked
        return getattr(self._inner, name)


class ThreadedWSGIServer(ThreadingMixIn, WSGIServer):
    daemon_threads = True


def read_only_middleware(app):
    """Reject mutating HTTP methods before they reach Optuna routes."""

    def wrapper(environ, start_response):
        method = environ.get("REQUEST_METHOD", "GET").upper()
        path = environ.get("PATH_INFO", "")
        if method not in _ALLOWED_METHODS:
            print(
                f"[read-only] blocked {method} {path}",
                file=sys.stderr,
                flush=True,
            )
            body = (
                b'{"reason":"read-only public dashboard: '
                b'write methods (POST/PUT/DELETE/PATCH) are disabled"}'
            )
            start_response(
                "403 Forbidden",
                [
                    ("Content-Type", "application/json"),
                    ("Content-Length", str(len(body))),
                    ("Cache-Control", "no-store"),
                    ("X-Optuna-Dashboard-Mode", "read-only"),
                ],
            )
            return [body]
        return app(environ, start_response)

    return wrapper


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description="Optuna Dashboard: public + read-only (HTTP + storage)."
    )
    parser.add_argument(
        "storage",
        help="Storage URL (e.g. sqlite:////path/to/presets.db)",
    )
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=8080)
    args = parser.parse_args(argv)

    inner = RDBStorage(args.storage)
    storage = ReadOnlyStorage(inner)
    app = read_only_middleware(wsgi(storage))

    print(
        f"=== READ-ONLY public dashboard pid={__import__('os').getpid()} ===",
        file=sys.stderr,
    )
    print(
        f"Listening on http://{args.host}:{args.port}/ "
        f"(HTTP 403 on writes + storage write lock)",
        file=sys.stderr,
    )
    print(
        "If DELETE study still works, you are NOT running this file "
        "(kill old optuna-dashboard and restart ./dashboard-public.sh).",
        file=sys.stderr,
    )
    httpd = make_server(args.host, args.port, app, server_class=ThreadedWSGIServer)
    try:
        httpd.serve_forever()
    except KeyboardInterrupt:
        print("\nStopped.", file=sys.stderr)
        return 0
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
