from __future__ import annotations

import os
from pathlib import Path

from dataikuapi import DSSClient


def require_env(name: str) -> str:
    value = os.getenv(name)
    if not value:
        raise RuntimeError(f"Missing required environment variable: {name}")
    return value


def find_plugin_archive(dist_dir: Path) -> Path:
    zip_files = sorted(dist_dir.glob("*.zip"))
    if not zip_files:
        raise FileNotFoundError(f"No plugin zip file found in {dist_dir}")
    if len(zip_files) > 1:
        print(f"Found multiple plugin archives, using the latest one: {zip_files[-1].name}")
    return zip_files[-1]


def parse_webapp_ids(raw_value: str | None) -> list[str]:
    if not raw_value:
        return []
    return [webapp_id.strip() for webapp_id in raw_value.split(",") if webapp_id.strip()]

def restart_playwright_webapps(client: DSSClient):
    print("Restarting playwright webapps ...")
    PLAYWRIGHT_PROJECT_KEY = "AH_PLAYWRIGHT"
    PLAYWRIGHT_WEBAPPS = {
        "q5ONPVi": "AH/playwright_test",
        "tUGUFDw": "AHL/playwright_test",
        "H3DKyAT": "AH/playwright_container"
        }
    playwright_project = client.get_project(PLAYWRIGHT_PROJECT_KEY)
    for webapp_id, webapp_label in PLAYWRIGHT_WEBAPPS.items():
        print(f"Restarting webapp '{webapp_label}' (id={webapp_id})...")
        webapp = playwright_project.get_webapp(webapp_id)
        webapp.start_or_restart_backend().wait_for_result()
        print(f"Webapp '{webapp_id}' has been restarted successfully!")
    print("✅ All webapps have been restarted successfully!")


def main() -> None:
    host = require_env("DSS_HOST")
    api_key = require_env("DSS_API_KEY")
    plugin_id = os.getenv("PLUGIN_ID", "agent-hub")

    client = DSSClient(host, api_key)
    plugin = client.get_plugin(plugin_id)
    print(f"Found plugin: {plugin_id}")

    archive_path = find_plugin_archive(Path("dist"))
    print(f"Using plugin archive: {archive_path}")

    with archive_path.open("rb") as file_obj:
        print("Uploading plugin zip ...")
        plugin.update_from_zip(file_obj)
        print("✅ Plugin updated.")

    print("Updating plugin code environment ...")
    plugin.update_code_env().wait_for_result()
    print("✅ Plugin code env updated.")

    restart_playwright_webapps(client)

    print("✅ Plugin update complete!")



if __name__ == "__main__":
    main()
