from datetime import datetime
from zoneinfo import ZoneInfo

from flask import Flask
from flask import jsonify
from flask import request
from flask import session
from flask_admin import Admin
from flask_admin.contrib.sqla import ModelView
from flask_admin.model import typefmt
from flask_sqlalchemy import SQLAlchemy
from markupsafe import Markup
from sqlalchemy import DateTime
from sqlalchemy import String
from sqlalchemy.orm import Mapped
from sqlalchemy.orm import mapped_column

app = Flask(__name__)
app.config["SECRET_KEY"] = "secret"
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///db.sqlite"
db = SQLAlchemy()
db.init_app(app)
admin = Admin(app, name="Example: Datetime and Timezone")


@app.route("/")
def index():
    return '<a href="/admin/timezone_aware_article">Click me to get to Admin!</a>'


@app.route("/set_timezone", methods=["POST"])
def set_timezone():
    """
    Save timezone to session so that datetime inputs can be correctly converted to UTC.
    """
    session.permanent = True
    timezone = request.get_json()
    if timezone:
        session["timezone"] = timezone
        return jsonify({"message": "Timezone set successfully"}), 200
    else:
        return jsonify({"error": "Invalid timezone"}), 400


class Article(db.Model):
    id: Mapped[int] = mapped_column(primary_key=True)
    text: Mapped[str] = mapped_column(String(30))
    last_edit: Mapped[datetime] = mapped_column(DateTime(timezone=True))


def date_format(view, value, name):
    """
    Ensure consistent date format and inject class for timezone.js parser.
    """
    if value is None:
        return ""
    return Markup(
        f'<span class="timezone-aware">{value.strftime("%Y-%m-%d %H:%M:%S")}</span>'
    )


MY_DEFAULT_FORMATTERS = dict(typefmt.BASE_FORMATTERS)
MY_DEFAULT_FORMATTERS.update(
    {
        datetime: date_format,
    }
)


class TimezoneAwareModelView(ModelView):
    column_type_formatters = MY_DEFAULT_FORMATTERS
    extra_js = ["/static/js/timezone.js"]

    def on_model_change(self, form, model, is_created):
        """
        Save datetime fields after converting from session['timezone'] to UTC.
        """
        user_timezone = session["timezone"]

        for field_name, field_value in form.data.items():
            if isinstance(field_value, datetime):
                # Convert naive datetime to timezone-aware datetime
                aware_time = field_value.replace(tzinfo=ZoneInfo(user_timezone))

                # Convert the time to UTC
                utc_time = aware_time.astimezone(ZoneInfo("UTC"))

                # Assign the UTC time to the model
                setattr(model, field_name, utc_time)

        super().on_model_change(form, model, is_created)


# inherit TimeZoneAwareModelView to make any admin page timezone-aware
class TimezoneAwareBlogModelView(TimezoneAwareModelView):
    column_labels = {
        "last_edit": "Last Edit (local time)",
    }


# compare with regular ModelView to display data as saved on db
class BlogModelView(ModelView):
    column_labels = {
        "last_edit": "Last Edit (UTC)",
    }


if __name__ == "__main__":
    with app.app_context():
        db.drop_all()
        db.create_all()
        db.session.add(
            Article(text="Written at 9:00 UTC", last_edit=datetime(2024, 8, 8, 9, 0, 0))
        )
        db.session.commit()
        admin.add_view(BlogModelView(Article, db, name="Article", endpoint="article"))
        admin.add_view(
            TimezoneAwareBlogModelView(
                Article,
                db,
                name="Timezone Aware Article",
                endpoint="timezone_aware_article",
            )
        )
    app.run(debug=True)
