"""Start a job with an Azure Function client and persist metadata in Azure Table."""

from __future__ import annotations

import asyncio
from collections.abc import Mapping
from dataclasses import dataclass
from datetime import datetime, timezone
from typing import Any

from infra_core.azure import AzureFunctionAsyncClient, AzureTableJobStore
from infra_core.azure.job_store import AzureTableRecordProtocol


@dataclass(frozen=True)
class JobRecord(AzureTableRecordProtocol):
    batch_id: str
    job_id: str
    status: str
    updated_at: str | None
    payload: dict[str, Any]

    def to_entity(self) -> dict[str, Any]:
        return {
            "PartitionKey": self.batch_id,
            "RowKey": self.job_id,
            "status": self.status,
            "updated_at": self.updated_at,
            "payload": self.payload,
        }

    @classmethod
    def from_entity(cls, entity: Mapping[str, Any]) -> JobRecord:
        return cls(
            batch_id=str(entity["PartitionKey"]),
            job_id=str(entity["RowKey"]),
            status=str(entity.get("status", "")),
            updated_at=entity.get("updated_at"),
            payload=dict(entity.get("payload", {})),
        )


async def main() -> None:
    client = AzureFunctionAsyncClient.from_env(
        base_url_env="MYAPP_AZURE_FUNCTION_BASE_URL",
        api_key_env="MYAPP_AZURE_FUNCTION_API_KEY",
    )
    response = await client.post_json("/jobs", json={"batch_id": "demo-batch"})
    await client.close()

    job_id = str(response["job_id"])
    with AzureTableJobStore[JobRecord](
        record_type=JobRecord,
        connection_string_var="AZURE_JOB_CONNECTION",
        table_name_var="AZURE_JOB_TABLE",
    ) as store:
        record = JobRecord(
            batch_id="demo-batch",
            job_id=job_id,
            status="submitted",
            updated_at=datetime.now(timezone.utc).isoformat(),
            payload=dict(response),
        )
        store.save(record)
        print(store.fetch("demo-batch", job_id))


if __name__ == "__main__":
    asyncio.run(main())
