"""AIGC 服务（与 npm AigcService 对齐）。"""
from typing import Any, Optional

from .base_client import BaseClient
from .types import ClientConfig, RequestOptions, TokenStorage


class AigcService(BaseClient):
    service_path = "aigc"

    def __init__(
        self,
        config: ClientConfig,
        token_storage: Optional[TokenStorage] = None,
    ) -> None:
        super().__init__(config, token_storage=token_storage)

    @property
    def service_name(self) -> str:
        return "aigc"

    def text_generation(
        self,
        request: dict,
        options: Optional[RequestOptions] = None,
    ) -> Any:
        return self.request("/text/generate", RequestOptions(method="POST", body=request))

    def image_generation(
        self,
        request: dict,
        options: Optional[RequestOptions] = None,
    ) -> Any:
        return self.request("/image/generate", RequestOptions(method="POST", body=request))

    def get_task_status(
        self,
        task_id: str,
        options: Optional[RequestOptions] = None,
    ) -> Any:
        return self.request(f"/task/{task_id}/status", RequestOptions(method="GET"))
