local_quantum_simulator/Dockerfile_build/source/modules/fasthtml.py

import asyncio
import logging
import os
from datetime import datetime, timedelta

import psutil
from connections.keycloak import get_valid_access_token, keycloak_openid
from connections.local_files import (
    delete_client_info,
    delete_token,
    delete_verifier,
    load_client_info,
    load_token,
    load_verifier,
    save_client_info,
    save_token,
    save_verifier,
)
from connections.quantum_backend import get_device_by_id, get_or_create_device_by_name
from connections.rabbitmq import (
    rabbitmq_manager,
)
from fasthtml.common import (
    H1,
    H2,
    H3,
    A,
    Button,
    Div,
    FastHTML,
    Form,
    Input,
    Meta,
    P,
    RedirectResponse,
    Script,
    Span,
    Strong,
    Titled,
)
from fasthtml.pico import Card, Container, picolink
from fasthtml.xtend import Style
from keycloak.pkce_utils import generate_code_challenge, generate_code_verifier
from modules.rabbitmq import consume_messages_topic

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


REDIRECT_URI = f"http://localhost:{os.environ['PORT']}/callback"

css = Style(":root {--pico-font-size:90%,--pico-font-family: Pacifico, cursive;}")
app = FastHTML(hdrs=(picolink, css))


@app.route("/", methods="get")
async def main_page():
    # Проверяем наличие валидного токена
    token_data = load_token()
    client_info = load_client_info()

    if token_data:
        expires_at = datetime.fromisoformat(token_data.get("expires_at", "2000-01-01"))
        if expires_at > datetime.now():
            # Если есть токен, но нет информации о клиенте, перенаправляем на регистрацию
            if not client_info:
                return RedirectResponse("/register-client", status_code=303)
            return RedirectResponse("/dashboard", status_code=303)
        else:
            delete_token()
            delete_client_info()

    return Titled(
        "Главная",
        H1("Добро пожаловать в Квантовый Симулятор"),
        P("Пожалуйста, войдите для продолжения"),
        A("Войти через Keycloak", href="/login", cls="button"),
    )


@app.route("/login", methods="get")
async def login():
    # Генерация PKCE кода верификатора и запроса
    code_verifier = generate_code_verifier()
    code_challenge, code_challenge_method = generate_code_challenge(
        code_verifier, method="S256"
    )

    # Сохраняем верификатор локально
    save_verifier(code_verifier)

    # Формирование URL авторизации
    auth_url = keycloak_openid.auth_url(
        redirect_uri=REDIRECT_URI,
        scope="openid profile email",
        code_challenge=code_challenge,
        code_challenge_method=code_challenge_method,
    )

    return RedirectResponse(auth_url, status_code=303)


@app.route("/callback", methods="get")
async def callback(request):
    code = request.query_params.get("code")

    if not code:
        return Titled(
            "Ошибка", P("Код не найден"), A("Попробовать снова", href="/login")
        )

    # Загружаем верификатор
    code_verifier = load_verifier()
    if not code_verifier:
        return Titled(
            "Ошибка",
            P("PKCE верификатор не найден"),
            A("Попробовать снова", href="/login"),
        )

    # Очищаем верификатор сразу после использования
    delete_verifier()

    try:
        # Обмен кода на токен
        token = keycloak_openid.token(
            grant_type="authorization_code",
            code=code,
            redirect_uri=REDIRECT_URI,
            code_verifier=code_verifier,
        )

        # Расчет времени истечения
        expires_in = token.get("expires_in", 3600)
        expires_at = datetime.now() + timedelta(seconds=expires_in)

        # Получение информации о пользователе
        access_token = token.get("access_token")
        if not access_token:
            raise Exception
        user_info = keycloak_openid.userinfo(access_token)

        # Сохраняем токен
        token_data_to_save = {
            "access_token": access_token,
            "refresh_token": token.get("refresh_token"),
            "expires_at": expires_at.isoformat(),
            "id_token": token.get("id_token"),
            "user_info": user_info,
        }

        save_token(token_data_to_save)

        # Перенаправление на регистрацию клиента (проверит, зарегистрирован ли уже)
        return RedirectResponse("/register-client", status_code=303)

    except Exception as e:
        logger.error(f"Ошибка при обмене токена: {e}")
        return Titled("Ошибка", P(str(e)), A("Попробовать снова", href="/login"))


@app.route("/register-client", methods="get")
async def register_client(request):
    """Страница регистрации клиента в центральном микросервисе"""
    token_data = load_token()

    if not token_data:
        return RedirectResponse("/login", status_code=303)

    # Проверяем, не зарегистрирован ли уже
    client_info = load_client_info()
    if client_info:
        return RedirectResponse("/dashboard", status_code=303)

    # Проверяем истечение токена
    expires_at = datetime.fromisoformat(str(token_data.get("expires_at")))
    if expires_at <= datetime.now():
        delete_token()
        return RedirectResponse("/login", status_code=303)

    user_info = token_data.get("user_info", {})
    username = user_info.get(
        "preferred_username", user_info.get("name", "Пользователь")
    )

    return Titled(
        "Регистрация клиента",
        Container(
            Card(
                H2("Добро пожаловать в Квантовый Симулятор!"),
                P(f"Здравствуйте, {username}!"),
                P("Это устройство должно быть зарегистрировано в центральном сервисе."),
                P(
                    "Пожалуйста, укажите уникальное имя и количество кубит для этого клиента:"
                ),
                Form(
                    Input(
                        type="text",
                        name="client_name",
                        id="client_name",
                        placeholder="например: рабочий-ноутбук, домашний-пк, raspberry-pi-01",
                        required=True,
                        style="width: 100%; padding: 10px; margin: 10px 0;",
                    ),
                    Input(
                        type="number",
                        name="number_qubits",
                        id="number_qubits",
                        placeholder=0,
                        required=True,
                        style="width: 100%; padding: 10px; margin: 10px 0;",
                    ),
                    Div(
                        Button(
                            "Зарегистрировать клиент",
                            type="submit",
                            cls="button",
                            style="background-color: #4CAF50;",
                        ),
                        style="margin-top: 10px;",
                    ),
                    hx_post="/register-client",
                    hx_target="#registration-result",
                    hx_swap="innerHTML",
                    style="margin-top: 20px;",
                ),
                Div(id="registration-result"),
            ),
            Div(
                Button(
                    "Выйти",
                    onclick="window.location.href='/logout'",
                    cls="button",
                    style="background-color: #f44336; margin-top: 20px;",
                ),
                style="text-align: center;",
            ),
        ),
    )


@app.route("/register-client", methods="post")
async def register_client_post(request):
    """Обработка POST запроса регистрации клиента"""
    form_data = await request.form()
    client_name = form_data.get("client_name")
    number_qubits = form_data.get("number_qubits")

    if not number_qubits:
        return Div(
            P("❌ Обязательно необходимо указать колчество кубит", style="color: red;"),
        )

    if not client_name:
        return Div(
            P("❌ Имя клиента обязательно", style="color: red;"),
        )

    # Проверка формата имени клиента
    import re

    if not re.match(r"^[a-zA-Z0-9_-]+$", client_name):
        return Div(
            P(
                "❌ Имя клиента может содержать только буквы, цифры, дефисы и подчеркивания",
                style="color: red;",
            ),
            Button("Попробовать снова", onclick="location.reload()", cls="button"),
        )

    token_data = load_token()
    if not token_data:
        return Div(
            P("❌ Сессия истекла. Пожалуйста, войдите снова.", style="color: red;"),
            A("Войти", href="/login", cls="button"),
        )

    access_token = token_data.get("access_token")

    # Регистрация в центральном микросервисе
    try:
        response = get_or_create_device_by_name(
            client_name, number_qubits, access_token
        )
        save_client_info(
            {
                "system_id": response["system_id"],
            }
        )

        return Div(
            P(
                "✅ Клиент успешно зарегистрирован!",
                style="color: green; font-weight: bold;",
            ),
            P(f"Имя клиента: {client_name}"),
            P("Перенаправление на панель управления..."),
            Meta(http_equiv="refresh", content="2;url=/dashboard"),
            Script("setTimeout(() => { window.location.href = '/dashboard'; }, 2000);"),
        )
    except Exception:
        return Div(
            P("❌ Ошибка регистрации", style="color: red;"),
            Button("Попробовать снова", onclick="location.reload()", cls="button"),
            P(
                "Если проблема повторяется, обратитесь к администратору.",
                style="font-size: 12px; margin-top: 10px;",
            ),
        )


@app.route("/dashboard", methods="get")
async def dashboard():
    access_token = get_valid_access_token()
    if rabbitmq_manager._connection and access_token:
        await rabbitmq_manager._connection.update_secret(
            access_token, reason="Token expired"
        )
    if not access_token:
        # Token refresh failed, redirect to login
        return RedirectResponse("/", status_code=303)
    token_data = load_token()
    client_info = load_client_info()

    if not token_data or not client_info:
        return RedirectResponse("/", status_code=303)

    user_info = token_data.get("user_info", {})
    username = user_info.get(
        "preferred_username", user_info.get("name", "Пользователь")
    )
    system_id = client_info.get("system_id", None)

    if not token_data:
        return Div(
            P("❌ Сессия истекла. Пожалуйста, войдите снова.", style="color: red;"),
            A("Войти", href="/login", cls="button"),
        )
    try:
        if system_id:
            device_data = get_device_by_id(system_id, access_token)
            device_name = device_data["system"]["system_name"]
            max_qubits = device_data["system"]["max_qubits"]
            created_at = device_data["system"]["created_at"]
        else:
            device_name = "error"
            max_qubits = "error"
            created_at = "error"
    except Exception as a:
        device_name = "error"
        max_qubits = "error"
        created_at = "error"

    return Titled(
        "Панель управления",
        Container(
            Card(
                P(Strong("Имя пользователя: "), username),
                P(Strong("Email: "), user_info.get("email", "Не указан")),
                P(Strong("Устройство зарегестрировано: "), created_at),
                Div(style="border: 1px solid black"),
                P(Strong("Имя вычислительной системы: "), device_name),
                P(Strong("Макс. количество кубит: "), max_qubits),
                Div(Button("Изменить", style="background-color: blue;")),
                Div(
                    id="client-status",
                    hx_get="/client-status",
                    hx_trigger="load, every 15s",
                ),
            ),
            Div(id="memory-stats", hx_get="/memory-stats", hx_trigger="load, every 5s"),
            A(
                Button(
                    "Выйти",
                    style="background-color: #f44336;",
                ),
                href="/logout",
                style="margin: 0px;color: inherit;text-decoration: inherit;",
            ),
        ),
    )


@app.route("/client-status", methods="get")
async def status():
    """Получить статус клиента из центрального микросервиса"""
    access_token = get_valid_access_token()
    client_info = load_client_info()

    if not access_token or not client_info:
        return P("Клиент не зарегистрирован", style="color: orange;")

    if rabbitmq_manager._connection:
        # TODO: FIX
        status = (
            rabbitmq_manager._connection.connected
        )  # get_client_status(client_name, access_token)

        if status:
            await consume_messages_topic()

            return Div(
                P(
                    Strong("Статус клиента: "),
                    Span("✅ Активен", style="color: green;"),
                ),
                style="padding: 10px; border-radius: 5px; margin: 10px 0;",
            )
        else:
            return Div(
                P(
                    Strong("Статус клиента: "),
                    Span("Не подключен", style="color: orange;"),
                ),
                P(
                    "Не удалось связаться с центральным сервисом или данному клиенту отказано в подключении"
                ),
                style="padding: 10px; border-radius: 5px; margin: 10px 0;",
            )
    else:
        return Div(
            P(
                Strong("Статус клиента: "),
                Span("Не подключен", style="color: orange;"),
            ),
            P(
                "Не удалось связаться с центральным сервисом или данному клиенту отказано в подключении"
            ),
            style="padding: 10px; border-radius: 5px; margin: 10px 0;",
        )


@app.route("/memory-stats", methods="get")
async def memory_stats():
    """Endpoint, возвращающий карточку использования памяти (для обновления через HTMX)"""
    process = psutil.Process(os.getpid())
    memory_info = process.memory_info()
    current_time = datetime.now()
    timestamp = current_time.strftime("%Y-%m-%d %H:%M:%S")

    def format_bytes(bytes):
        for unit in ["Б", "КБ", "МБ", "ГБ"]:
            if bytes < 1024.0:
                return f"{bytes:.1f} {unit}"
            bytes /= 1024.0
        return f"{bytes:.1f} ГБ"

    mem = 0
    if os.path.isfile("/sys/fs/cgroup/memory.max"):
        with open("/sys/fs/cgroup/memory.max") as limit:
            try:
                mem = int(limit.read())
            except Exception:
                mem = 0
            # resource.setrlimit(resource.RLIMIT_AS, (mem, mem))

    rss_memory = format_bytes(memory_info.rss)
    if mem > 0:
        memory_percent = (memory_info.rss / mem) * 100
    else:
        memory_percent = process.memory_percent()

    bar_color = (
        "#4CAF50"
        if memory_percent < 1
        else "#FFC107"
        if memory_percent < 5
        else "#F44336"
    )

    return Div(
        H2("Использование памяти"),
        Card(
            Div(
                H3(rss_memory),
                style="width: 100%;display:flex;flex-direction:column;align-items:center",
            ),
            Div(
                Div(
                    style=f"width: {min(memory_percent, 100)}%; background-color: {bar_color}; height: 20px; border-radius: 10px;"
                ),
                style="width: 100%; background-color: #e0e0e0; border-radius: 10px; overflow: hidden; margin: 10px 0;",
            ),
            P(
                f"{memory_percent:.2f}% от системной памяти",
                style="text-align: center; font-size: 12px;",
            ),
            P(
                Strong("Обновлено: "),
                timestamp,
                style="text-align: center; font-size: 11px; color: #666;",
            ),
        ),
    )


@app.route("/logout", methods="get")
async def logout():
    token_data = load_token()

    if token_data and token_data.get("refresh_token"):
        keycloak_openid.logout(token_data["refresh_token"])
    # Очищаем локальные файлы
    delete_token()
    delete_client_info()
    delete_verifier()

    return Titled(
        "Выход выполнен",
        H2("До свидания!"),
        P("Вы успешно вышли из системы."),
        P("Ваш клиент был отменен в центральном сервисе."),
        A("Войти снова", href="/", cls="button"),
    )


# Опционально: endpoint для проверки работоспособности для центрального сервиса
@app.route("/health", methods="get")
async def health():
    return {"status": "healthy", "timestamp": datetime.now().isoformat()}