import json
from collections.abc import Callable
from typing import Optional

from flask import url_for
from flask.sansio.app import App
from pydantic import BaseModel, RootModel

from app.config import Environment


class Asset(BaseModel):
    file: str
    name: str | None = None

    # Shared chunks emitted by vite (modules imported by more than one entrypoint) have no `src` or `isEntry`.
    src: str | None = None
    isEntry: bool = False


class Manifest(RootModel[Optional[dict[str, Asset]]]):
    root: dict[str, Asset] | None

    def __getitem__(self, item: str) -> Asset | None:
        if self.root:
            try:
                return self.root[item]
            except KeyError:
                pass
        return None


class FlaskAssetsViteExtension:
    def __init__(self) -> None:
        self._app: App | None = None
        self._live_enabled: bool | None = None
        self._manifest: Manifest | None = None

    def init_app(self, app: App) -> None:
        self._app = app
        self._live_enabled = self._app.config["ASSETS_VITE_LIVE_ENABLED"]

        try:
            with open("app/assets/dist/manifest.json") as f:
                data = json.load(f)
                self._manifest = Manifest(root=data)
        except FileNotFoundError as e:
            if not self._live_enabled or self._app.config["FLASK_ENV"] == Environment.PROD:
                raise Exception("Asset manifest not found, make sure assets are generated") from e

        app.extensions["flask_assets_vite"] = self
        app.context_processor(self.assets_processor)

    def assets_processor(self) -> dict[str, Callable[[str], str]]:
        def vite_asset(relative_file_path: str, *, external: bool | None = None) -> str:
            """
            Point assets at a vite development server while running locally
            to enable hot module replacement and automatic udpates to both SCSS
            and JavaScript.
            """
            if self._app is None or self._live_enabled is None:
                raise RuntimeError("Need to call FlaskAssetsViteExtension.init_app() first")

            if self._live_enabled:
                return f"{self._app.config['ASSETS_VITE_BASE_URL']}/static/{relative_file_path}"

            if not self._manifest:
                raise Exception("Asset manifest not found, make sure assets are generated")

            generated_asset = self._manifest[relative_file_path]
            if generated_asset:
                # assets that have been transpiled by vite should reference their hashed names
                return url_for("static", filename=generated_asset.file, _external=external)

            # assets that have not been transpiled by vite but may have been copied
            return url_for("static", filename=relative_file_path, _external=external)

        return dict(vite_asset=vite_asset)
