vmware-api-simulator/app/vsphere/domain/content.py

"""Content library + datastore file metadata."""

from __future__ import annotations

import secrets
from typing import Any

from asyncpg.exceptions import UniqueViolationError  # type: ignore[import-untyped]

from app.db.pool import Database
from app.vsphere import inventory
from app.vsphere.domain import tagging
from app.vsphere.domain import tasks as task_store
from app.vsphere.errors import already_exists, invalid_argument, not_found


def _pool(database: Database) -> Any:
    return database.pool  # type: ignore[attr-defined]


async def create_library(
    database: Database,
    *,
    name: str,
    description: str = "",
    library_id: str | None = None,
) -> str:
    if not name.strip():
        raise invalid_argument("name is required")
    lib_id = library_id or f"lib-{secrets.token_hex(6)}"
    pool = _pool(database)
    try:
        async with pool.acquire() as conn:
            await conn.execute(
                """
                INSERT INTO vsphere_libraries (id, name, description, type, props)
                VALUES ($1, $2, $3, 'LOCAL', '{}'::jsonb)
                ON CONFLICT (id) DO NOTHING
                """,
                lib_id,
                name,
                description,
            )
            # Id already present (seed ensure) — treat as success.
            exists = await conn.fetchval("SELECT 1 FROM vsphere_libraries WHERE id = $1", lib_id)
            if exists:
                return lib_id
    except UniqueViolationError as error:
        # Name collision with a different id: keep requesting id with a unique name.
        if library_id:
            pool = _pool(database)
            async with pool.acquire() as conn:
                await conn.execute(
                    """
                    INSERT INTO vsphere_libraries (id, name, description, type, props)
                    VALUES ($1, $2, $3, 'LOCAL', '{}'::jsonb)
                    ON CONFLICT (id) DO NOTHING
                    """,
                    lib_id,
                    f"{name} ({lib_id})",
                    description,
                )
            return lib_id
        raise already_exists(f"Library {name} already exists") from error
    return lib_id


async def list_libraries(database: Database) -> list[dict[str, Any]]:
    pool = _pool(database)
    async with pool.acquire() as conn:
        rows = await conn.fetch("SELECT * FROM vsphere_libraries ORDER BY name")
    return [
        {
            "id": row["id"],
            "name": row["name"],
            "description": row["description"],
            "type": row["type"],
        }
        for row in rows
    ]


async def create_library_item(
    database: Database,
    *,
    library_id: str,
    name: str,
    item_type: str = "ovf",
    description: str = "",
    item_id: str | None = None,
) -> str:
    pool = _pool(database)
    async with pool.acquire() as conn:
        exists = await conn.fetchval("SELECT 1 FROM vsphere_libraries WHERE id = $1", library_id)
        if not exists:
            raise not_found(f"Library {library_id} not found")
        resolved_id = item_id or f"item-{secrets.token_hex(6)}"
        try:
            await conn.execute(
                """
                INSERT INTO vsphere_library_items (id, library_id, name, type, description)
                VALUES ($1, $2, $3, $4, $5)
                ON CONFLICT (id) DO NOTHING
                """,
                resolved_id,
                library_id,
                name,
                item_type,
                description,
            )
        except UniqueViolationError:
            # (library_id, name) taken — keep stable id with a unique name for seed.
            if item_id:
                await conn.execute(
                    """
                    INSERT INTO vsphere_library_items (id, library_id, name, type, description)
                    VALUES ($1, $2, $3, $4, $5)
                    ON CONFLICT (id) DO NOTHING
                    """,
                    resolved_id,
                    library_id,
                    f"{name}-{resolved_id}",
                    item_type,
                    description,
                )
            else:
                raise
    return resolved_id


async def list_library_items(database: Database, library_id: str) -> list[dict[str, Any]]:
    pool = _pool(database)
    async with pool.acquire() as conn:
        exists = await conn.fetchval("SELECT 1 FROM vsphere_libraries WHERE id = $1", library_id)
        if not exists:
            raise not_found(f"Library {library_id} not found")
        rows = await conn.fetch(
            "SELECT * FROM vsphere_library_items WHERE library_id = $1 ORDER BY name",
            library_id,
        )
    return [
        {
            "id": row["id"],
            "library_id": row["library_id"],
            "name": row["name"],
            "type": row["type"],
            "description": row["description"],
        }
        for row in rows
    ]


_LAB_SESSION_ID = "session-lab-1"
_LAB_ITEM_ID = "item-ubuntu"


async def clear_transfer_sessions(database: Database) -> None:
    """Drop durable transfer sessions (used by inventory reseed)."""

    pool = _pool(database)
    async with pool.acquire() as conn:
        await conn.execute(
            """
            DO $$ BEGIN
              IF to_regclass('public.vsphere_transfer_sessions') IS NOT NULL THEN
                DELETE FROM vsphere_transfer_sessions;
              END IF;
            END $$;
            """
        )


def _decode_files(value: Any) -> dict[str, Any]:
    import json

    current = value
    while isinstance(current, str):
        current = json.loads(current)
    if isinstance(current, dict):
        return current
    return {}


async def _upsert_session(
    database: Database,
    *,
    session_id: str,
    kind: str,
    library_item_id: str,
    state: str,
    files: dict[str, Any],
) -> None:
    import json

    pool = _pool(database)
    async with pool.acquire() as conn:
        await conn.execute(
            """
            INSERT INTO vsphere_transfer_sessions (id, kind, library_item_id, state, files, updated_at)
            VALUES ($1, $2, $3, $4, $5::jsonb, now())
            ON CONFLICT (id, kind) DO UPDATE SET
              library_item_id = EXCLUDED.library_item_id,
              state = EXCLUDED.state,
              files = EXCLUDED.files,
              updated_at = now()
            """,
            session_id,
            kind,
            library_item_id,
            state,
            json.dumps(files),
        )


async def _get_session_row(
    database: Database, session_id: str, kind: str | None = None
) -> dict[str, Any] | None:
    pool = _pool(database)
    async with pool.acquire() as conn:
        if kind:
            row = await conn.fetchrow(
                "SELECT * FROM vsphere_transfer_sessions WHERE id = $1 AND kind = $2",
                session_id,
                kind,
            )
        else:
            row = await conn.fetchrow(
                "SELECT * FROM vsphere_transfer_sessions WHERE id = $1",
                session_id,
            )
    if row is None:
        return None
    return {
        "id": row["id"],
        "kind": row["kind"],
        "library_item_id": row["library_item_id"],
        "state": row["state"],
        "files": _decode_files(row["files"]),
    }


async def create_update_session(
    database: Database,
    *,
    library_item_id: str,
    session_id: str | None = None,
) -> str:
    pool = _pool(database)
    async with pool.acquire() as conn:
        exists = await conn.fetchval(
            "SELECT 1 FROM vsphere_library_items WHERE id = $1", library_item_id
        )
    if not exists:
        raise not_found(f"Library item {library_item_id} not found")
    resolved = session_id or f"update-session-{secrets.token_hex(6)}"
    files = {
        "upload.bin": {
            "name": "upload.bin",
            "source_type": "PUSH",
            "size": 1024,
            "status": "READY",
            "upload_endpoint": {
                "uri": f"/api/content/library/item/update-session/{resolved}/file/upload.bin",
            },
        }
    }
    await _upsert_session(
        database,
        session_id=resolved,
        kind="update",
        library_item_id=library_item_id,
        state="ACTIVE",
        files=files,
    )
    return resolved


async def add_update_session_file(
    database: Database,
    session_id: str,
    *,
    name: str,
    source_type: str = "PUSH",
    size: int = 0,
    content: str = "",
) -> dict[str, Any]:
    session = await _get_session_row(database, session_id, "update")
    if session is None:
        raise not_found(f"Update session {session_id} not found")
    file_info = {
        "name": name,
        "source_type": source_type,
        "size": size or len(content.encode("utf-8")),
        "status": "READY",
        "upload_endpoint": {
            "uri": f"/api/content/library/item/update-session/{session_id}/file/{name}",
        },
        "content": content,
    }
    files = dict(session["files"])
    files[name] = file_info
    await _upsert_session(
        database,
        session_id=session_id,
        kind="update",
        library_item_id=session["library_item_id"],
        state=session["state"],
        files=files,
    )
    return {
        "name": name,
        "source_type": source_type,
        "size": file_info["size"],
        "status": "READY",
        "upload_endpoint": file_info["upload_endpoint"],
    }


async def complete_update_session(database: Database, session_id: str) -> None:
    session = await _get_session_row(database, session_id, "update")
    if session is None:
        raise not_found(f"Update session {session_id} not found")
    await _upsert_session(
        database,
        session_id=session_id,
        kind="update",
        library_item_id=session["library_item_id"],
        state="DONE",
        files=session["files"],
    )


async def get_update_session(database: Database, session_id: str) -> dict[str, Any]:
    session = await _get_session_row(database, session_id, "update")
    if session is None:
        raise not_found(f"Update session {session_id} not found")
    return {
        "id": session["id"],
        "library_item_id": session["library_item_id"],
        "state": session["state"],
        "client_progress": 100 if session["state"] == "DONE" else 50,
    }


async def create_download_session(
    database: Database,
    *,
    library_item_id: str,
    session_id: str | None = None,
) -> str:
    pool = _pool(database)
    async with pool.acquire() as conn:
        exists = await conn.fetchval(
            "SELECT 1 FROM vsphere_library_items WHERE id = $1", library_item_id
        )
    if not exists:
        raise not_found(f"Library item {library_item_id} not found")
    resolved = session_id or f"download-session-{secrets.token_hex(6)}"
    files = {
        "descriptor.ovf": {
            "name": "descriptor.ovf",
            "size": 256,
            "status": "READY",
            "download_endpoint": {
                "uri": f"/api/content/library/item/download-session/{resolved}/file/descriptor.ovf",
            },
        }
    }
    await _upsert_session(
        database,
        session_id=resolved,
        kind="download",
        library_item_id=library_item_id,
        state="ACTIVE",
        files=files,
    )
    return resolved


async def get_download_session(database: Database, session_id: str) -> dict[str, Any]:
    session = await _get_session_row(database, session_id, "download")
    if session is None:
        raise not_found(f"Download session {session_id} not found")
    return {
        "id": session["id"],
        "library_item_id": session["library_item_id"],
        "state": session["state"],
    }


async def list_download_session_files(database: Database, session_id: str) -> list[dict[str, Any]]:
    session = await _get_session_row(database, session_id, "download")
    if session is None:
        raise not_found(f"Download session {session_id} not found")
    return [
        {
            "name": f["name"],
            "size": f["size"],
            "status": f["status"],
            "download_endpoint": f["download_endpoint"],
        }
        for f in session["files"].values()
    ]


async def deploy_ovf_from_library(
    database: Database,
    *,
    item_id: str,
    name: str,
    folder: str = "group-v23",
    host: str = "host-11",
    datastore: str = "datastore-31",
) -> tuple[str, str]:
    pool = _pool(database)
    async with pool.acquire() as conn:
        row = await conn.fetchrow("SELECT * FROM vsphere_library_items WHERE id = $1", item_id)
    if row is None:
        raise not_found(f"Library item {item_id} not found")
    moid = await inventory.next_moid(database, "vm")
    await inventory.upsert_object(
        database,
        moid=moid,
        type_name="VirtualMachine",
        name=name,
        parent_moid=folder,
        props={
            "power_state": "POWERED_OFF",
            "cpu_count": 2,
            "memory_size_mib": 2048,
            "guest_OS": "OTHER_GUEST_64",
            "hardware_version": "VMX_19",
            "host": host,
            "datastore": datastore,
            "networks": ["network-41"],
            "identity": {"name": name},
            "deployed_from_library_item": item_id,
            "nics": [],
            "disks": [
                {
                    "key": "2000",
                    "value": {"label": "Hard disk 1", "capacity": 21474836480, "type": "SCSI"},
                }
            ],
        },
    )
    task_id = await task_store.create_task(
        database,
        description=f"Deploy OVF item {item_id} as {moid}",
        service="com.vmware.vcenter.ovf",
        operation="deploy",
        result={"vm": moid},
    )
    return moid, task_id


async def list_datastore_files(database: Database, datastore: str) -> list[dict[str, Any]]:
    obj = await inventory.get_object(database, datastore)
    if obj is None or obj.type != "Datastore":
        raise not_found(f"Datastore {datastore} not found")
    pool = _pool(database)
    async with pool.acquire() as conn:
        rows = await conn.fetch(
            """
            SELECT path, size, type FROM vsphere_datastore_files
            WHERE datastore_moid = $1 ORDER BY path
            """,
            datastore,
        )
    return [{"path": row["path"], "size": row["size"], "type": row["type"]} for row in rows]


async def put_datastore_file(
    database: Database,
    datastore: str,
    path: str,
    *,
    size: int = 0,
    file_type: str = "FILE",
) -> None:
    obj = await inventory.get_object(database, datastore)
    if obj is None or obj.type != "Datastore":
        raise not_found(f"Datastore {datastore} not found")
    pool = _pool(database)
    async with pool.acquire() as conn:
        await conn.execute(
            """
            INSERT INTO vsphere_datastore_files (datastore_moid, path, size, type)
            VALUES ($1, $2, $3, $4)
            ON CONFLICT (datastore_moid, path) DO UPDATE SET size = EXCLUDED.size, type = EXCLUDED.type
            """,
            datastore,
            path,
            size,
            file_type,
        )


async def seed_platform_extras(database: Database) -> None:
    """Idempotent demo libraries/tags/files/sessions when platform tables exist.

    Always ensures stable lab IDs so probes and clients hit non-empty GETs:
    ``lib-local-1``, ``item-ubuntu``, ``cat-lab-1``, ``tag-lab-1``, ``session-lab-1``, …
    """

    from app.vsphere.domain import tasks as task_store

    # Ensure lab libraries/items even when older random-id rows already exist.
    lib = await create_library(
        database,
        name="Local Content",
        description="Lab library",
        library_id="lib-local-1",
    )
    await create_library_item(
        database,
        library_id=lib,
        name="ubuntu-22.04",
        item_type="ovf",
        description="Ubuntu template OVF",
        item_id="item-ubuntu",
    )
    await create_library_item(
        database,
        library_id=lib,
        name="centos-stream-9",
        item_type="ovf",
        description="CentOS Stream OVF",
        item_id="item-centos",
    )
    pub = await create_library(
        database,
        name="Published Templates",
        description="Published lab library",
        library_id="lib-published-1",
    )
    await create_library_item(
        database,
        library_id=pub,
        name="golden-image",
        item_type="ovf",
        description="Golden image OVF",
        item_id="item-golden",
    )

    env = await tagging.create_category(
        database,
        name="Environment",
        description="Environment tags",
        associable_types=["VirtualMachine", "HostSystem"],
        category_id="urn:vmomi:InventoryServiceCategory:environment:GLOBAL",
    )
    owner = await tagging.create_category(
        database,
        name="Owner",
        description="Team ownership",
        associable_types=["VirtualMachine"],
        category_id="urn:vmomi:InventoryServiceCategory:owner:GLOBAL",
    )
    prod = await tagging.create_tag(
        database,
        category_id=env,
        name="prod",
        tag_id="urn:vmomi:InventoryServiceTag:prod:GLOBAL",
    )
    await tagging.create_tag(
        database,
        category_id=env,
        name="staging",
        tag_id="urn:vmomi:InventoryServiceTag:staging:GLOBAL",
    )
    await tagging.create_tag(
        database,
        category_id=owner,
        name="platform",
        tag_id="urn:vmomi:InventoryServiceTag:platform:GLOBAL",
    )
    await tagging.create_category(
        database,
        name="Lab",
        description="Probe alias category",
        associable_types=["VirtualMachine"],
        category_id="cat-lab-1",
    )
    await tagging.create_tag(
        database,
        category_id="cat-lab-1",
        name="lab",
        tag_id="tag-lab-1",
    )
    await tagging.attach_tag(database, prod, "VirtualMachine", "vm-101")
    await tagging.attach_tag(database, prod, "VirtualMachine", "vm-102")
    await tagging.attach_tag(database, "tag-lab-1", "VirtualMachine", "vm-101")

    await put_datastore_file(
        database, "datastore-31", "[datastore1] ISO/ubuntu.iso", size=900000000
    )
    await put_datastore_file(
        database, "datastore-31", "[datastore1] ISO/vmware-tools.iso", size=120000000
    )
    await put_datastore_file(database, "datastore-31", "[datastore1] web-01/web-01.vmx", size=4096)
    await put_datastore_file(
        database, "datastore-31", "[datastore1] web-01/web-01.vmdk", size=42949672960
    )

    item_id = _LAB_ITEM_ID
    pool = _pool(database)
    async with pool.acquire() as conn:
        has_item = await conn.fetchval("SELECT 1 FROM vsphere_library_items WHERE id = $1", item_id)
    if has_item:
        await create_download_session(database, library_item_id=item_id, session_id=_LAB_SESSION_ID)
        await create_update_session(database, library_item_id=item_id, session_id=_LAB_SESSION_ID)

    # Lab snapshots so GET /snapshots is DB-backed (no create-on-read).
    from app.vsphere.domain import vm_ops

    for vm_moid in ("vm-101", "vm-102"):
        snaps = await vm_ops.list_snapshots(database, vm_moid)
        if not snaps:
            await vm_ops.create_snapshot(
                database, vm_moid, name="initial", description="Lab default snapshot"
            )

    await task_store.create_task(
        database,
        description="Lab inventory seed",
        service="com.vmware.vcenter",
        operation="seed",
        status="SUCCEEDED",
        result={"seeded": True, "profile": "lab"},
        task_id="task-1",
    )