flexmeasures-main/flexmeasures/ui/views/assets/forms.py

from __future__ import annotations

import copy

from flask import current_app
from flask_security import current_user
from flask_wtf import FlaskForm
from sqlalchemy import select
from wtforms import StringField, DecimalField, SelectField, IntegerField
from wtforms.validators import DataRequired, optional

from flexmeasures.auth.policy import user_has_admin_access
from flexmeasures.data import db
from flexmeasures.data.models.generic_assets import GenericAssetType
from flexmeasures.data.models.user import Account


class AssetForm(FlaskForm):
    """The default asset form only allows to edit the name and location."""

    name = StringField("Name")
    latitude = DecimalField(
        "Latitude",
        validators=[optional()],
        places=None,
        render_kw={"placeholder": "--Click the map or enter a latitude--"},
    )
    longitude = DecimalField(
        "Longitude",
        validators=[optional()],
        places=None,
        render_kw={"placeholder": "--Click the map or enter a longitude--"},
    )
    attributes = StringField("Other attributes (JSON)", default="{}")

    def validate_on_submit(self):
        if (
            hasattr(self, "generic_asset_type_id")
            and self.generic_asset_type_id.data == -1
        ):
            self.generic_asset_type_id.data = (
                ""  # cannot be coerced to int so will be flagged as invalid input
            )
        if hasattr(self, "account_id") and self.account_id.data == -1:
            del self.account_id  # asset will be public
        result = super().validate_on_submit()
        return result

    def to_json(self) -> dict:
        """turn form data into a JSON we can POST to our internal API"""
        data = copy.copy(self.data)
        if data.get("longitude") is not None:
            data["longitude"] = float(data["longitude"])
        if data.get("latitude") is not None:
            data["latitude"] = float(data["latitude"])
        if data.get("parent_asset_id") is not None:
            data["parent_asset_id"] = int(data["parent_asset_id"])
        if "csrf_token" in data:
            del data["csrf_token"]

        return data

    def process_api_validation_errors(self, api_response: dict):
        """Process form errors from the API for the WTForm"""
        if not isinstance(api_response, dict):
            return
        for error_header in ("json", "validation_errors"):
            if error_header not in api_response:
                continue
            for field in list(self._fields.keys()):
                if field in list(api_response[error_header].keys()):
                    field_errors = api_response[error_header][field]
                    if isinstance(field_errors, list):
                        self._fields[field].errors += api_response[error_header][field]
                    else:
                        self._fields[field].errors.append(
                            api_response[error_header][field]
                        )

    def with_options(self):
        if "generic_asset_type_id" in self:
            self.generic_asset_type_id.choices = [(-1, "--Select type--")] + [
                (atype.id, atype.name)
                for atype in db.session.scalars(select(GenericAssetType)).all()
            ]
        if "account_id" in self:
            self.account_id.choices = [(-1, "--Select account--")] + [
                (account.id, account.name)
                for account in db.session.scalars(select(Account)).all()
            ]


class NewAssetForm(AssetForm):
    """Here, in addition, we allow to set asset type and account."""

    generic_asset_type_id = SelectField(
        "Asset type", coerce=int, validators=[DataRequired()]
    )
    account_id = SelectField("Account", coerce=int)
    parent_asset_id = IntegerField(
        "Parent Asset Id", validators=[optional()]
    )  # Add parent_id field

    def set_account(self) -> tuple[Account | None, str | None]:
        """Set an account for the to-be-created asset.
        Return the account (if available) and an error message"""

        account_error = None

        if self.account_id.data == -1:
            if user_has_admin_access(current_user, "update"):
                return None, None  # Account can be None (public asset)
            else:
                account_error = "Please pick an existing account."

        account = db.session.execute(
            select(Account).filter_by(id=int(self.account_id.data))
        ).scalar_one_or_none()

        if account:
            self.account_id.data = account.id
        else:
            current_app.logger.error(account_error)
        return account, account_error

    def set_asset_type(self) -> tuple[GenericAssetType | None, str | None]:
        """Set an asset type for the to-be-created asset.
        Return the asset type (if available) and an error message."""
        asset_type = None
        asset_type_error = None

        if int(self.generic_asset_type_id.data) == -1:
            asset_type_error = "Pick an existing asset type."
        else:
            asset_type = db.session.execute(
                select(GenericAssetType).filter_by(
                    id=int(self.generic_asset_type_id.data)
                )
            ).scalar_one_or_none()

        if asset_type:
            self.generic_asset_type_id.data = asset_type.id
        else:
            current_app.logger.error(asset_type_error)
        return asset_type, asset_type_error