Byted Mediakit Process Tools

#!/bin/bash
# AMK音视频处理工具集启动脚本
# 兼容各种执行路径
SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" &>/dev/null && pwd)
cd "$SCRIPT_DIR/scripts" || exit 1
if [ -f ".venv/bin/activate" ]; then
  source .venv/bin/activate
fi
if command -v python3 >/dev/null 2>&1; then
  python3 main.py "$@"
else
  python main.py "$@"
fi

                                 Apache License
                           Version 2.0, January 2004
                        http://www.apache.org/licenses/

   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION

   1. Definitions.

      "License" shall mean the terms and conditions for use, reproduction,
      and distribution as defined by Sections 1 through 9 of this document.

      "Licensor" shall mean the copyright owner or entity authorized by
      the copyright owner that is granting the License.

      "Legal Entity" shall mean the union of the acting entity and all
      other entities that control, are controlled by, or are under common
      control with that entity. For the purposes of this definition,
      "control" means (i) the power, direct or indirect, to cause the
      direction or management of such entity, whether by contract or
      otherwise, or (ii) ownership of fifty percent (50%) or more of the
      outstanding shares, or (iii) beneficial ownership of such entity.

      "You" (or "Your") shall mean an individual or Legal Entity
      exercising permissions granted by this License.

      "Source" form shall mean the preferred form for making modifications,
      including but not limited to software source code, documentation
      source, and configuration files.

      "Object" form shall mean any form resulting from mechanical
      transformation or translation of a Source form, including but
      not limited to compiled object code, generated documentation,
      and conversions to other media types.

      "Work" shall mean the work of authorship, whether in Source or
      Object form, made available under the License, as indicated by a
      copyright notice that is included in or attached to the work
      (an example is provided in the Appendix below).

      "Derivative Works" shall mean any work, whether in Source or Object
      form, that is based on (or derived from) the Work and for which the
      editorial revisions, annotations, elaborations, or other modifications
      represent, as a whole, an original work of authorship. For the purposes
      of this License, Derivative Works shall not include works that remain
      separable from, or merely link (or bind by name) to the interfaces of,
      the Work and Derivative Works thereof.

      "Contribution" shall mean any work of authorship, including
      the original version of the Work and any modifications or additions
      to that Work or Derivative Works thereof, that is intentionally
      submitted to Licensor for inclusion in the Work by the copyright owner
      or by an individual or Legal Entity authorized to submit on behalf of
      the copyright owner. For the purposes of this definition, "submitted"
      means any form of electronic, verbal, or written communication sent
      to the Licensor or its representatives, including but not limited to
      communication on electronic mailing lists, source code control systems,
      and issue tracking systems that are managed by, or on behalf of, the
      Licensor for the purpose of discussing and improving the Work, but
      excluding communication that is conspicuously marked or otherwise
      designated in writing by the copyright owner as "Not a Contribution."

      "Contributor" shall mean Licensor and any individual or Legal Entity
      on behalf of whom a Contribution has been received by Licensor and
      subsequently incorporated within the Work.

   2. Grant of Copyright License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      copyright license to reproduce, prepare Derivative Works of,
      publicly display, publicly perform, sublicense, and distribute the
      Work and such Derivative Works in Source or Object form.

   3. Grant of Patent License. Subject to the terms and conditions of
      this License, each Contributor hereby grants to You a perpetual,
      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
      (except as stated in this section) patent license to make, have made,
      use, offer to sell, sell, import, and otherwise transfer the Work,
      where such license applies only to those patent claims licensable
      by such Contributor that are necessarily infringed by their
      Contribution(s) alone or by combination of their Contribution(s)
      with the Work to which such Contribution(s) was submitted. If You
      institute patent litigation against any entity (including a
      cross-claim or counterclaim in a lawsuit) alleging that the Work
      or a Contribution incorporated within the Work constitutes direct
      or contributory patent infringement, then any patent licenses
      granted to You under this License for that Work shall terminate
      as of the date such litigation is filed.

   4. Redistribution. You may reproduce and distribute copies of the
      Work or Derivative Works thereof in any medium, with or without
      modifications, and in Source or Object form, provided that You
      meet the following conditions:

      (a) You must give any other recipients of the Work or
          Derivative Works a copy of this License; and

      (b) You must cause any modified files to carry prominent notices
          stating that You changed the files; and

      (c) You must retain, in the Source form of any Derivative Works
          that You distribute, all copyright, patent, trademark, and
          attribution notices from the Source form of the Work,
          excluding those notices that do not pertain to any part of
          the Derivative Works; and

      (d) If the Work includes a "NOTICE" text file as part of its
          distribution, then any Derivative Works that You distribute must
          include a readable copy of the attribution notices contained
          within such NOTICE file, excluding those notices that do not
          pertain to any part of the Derivative Works, in at least one
          of the following places: within a NOTICE text file distributed
          as part of the Derivative Works; within the Source form or
          documentation, if provided along with the Derivative Works; or,
          within a display generated by the Derivative Works, if and
          wherever such third-party notices normally appear. The contents
          of the NOTICE file are for informational purposes only and
          do not modify the License. You may add Your own attribution
          notices within Derivative Works that You distribute, alongside
          or as an addendum to the NOTICE text from the Work, provided
          that such additional attribution notices cannot be construed
          as modifying the License.

      You may add Your own copyright statement to Your modifications and
      may provide additional or different license terms and conditions
      for use, reproduction, or distribution of Your modifications, or
      for any such Derivative Works as a whole, provided Your use,
      reproduction, and distribution of the Work otherwise complies with
      the conditions stated in this License.

   5. Submission of Contributions. Unless You explicitly state otherwise,
      any Contribution intentionally submitted for inclusion in the Work
      by You to the Licensor shall be under the terms and conditions of
      this License, without any additional terms or conditions.
      Notwithstanding the above, nothing herein shall supersede or modify
      the terms of any separate license agreement you may have executed
      with Licensor regarding such Contributions.

   6. Trademarks. This License does not grant permission to use the trade
      names, trademarks, service marks, or product names of the Licensor,
      except as required for reasonable and customary use in describing the
      origin of the Work and reproducing the content of the NOTICE file.

   7. Disclaimer of Warranty. Unless required by applicable law or
      agreed to in writing, Licensor provides the Work (and each
      Contributor provides its Contributions) on an "AS IS" BASIS,
      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
      implied, including, without limitation, any warranties or conditions
      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
      PARTICULAR PURPOSE. You are solely responsible for determining the
      appropriateness of using or redistributing the Work and assume any
      risks associated with Your exercise of permissions under this License.

   8. Limitation of Liability. In no event and under no legal theory,
      whether in tort (including negligence), contract, or otherwise,
      unless required by applicable law (such as deliberate and grossly
      negligent acts) or agreed to in writing, shall any Contributor be
      liable to You for damages, including any direct, indirect, special,
      incidental, or consequential damages of any character arising as a
      result of this License or out of the use or inability to use the
      Work (including but not limited to damages for loss of goodwill,
      work stoppage, computer failure or malfunction, or any and all
      other commercial damages or losses), even if such Contributor
      has been advised of the possibility of such damages.

   9. Accepting Warranty or Additional Liability. While redistributing
      the Work or Derivative Works thereof, You may choose to offer,
      and charge a fee for, acceptance of support, warranty, indemnity,
      or other liability obligations and/or rights consistent with this
      License. However, in accepting such obligations, You may act only
      on Your own behalf and on Your sole responsibility, not on behalf
      of any other Contributor, and only if You agree to indemnify,
      defend, and hold each Contributor harmless for any liability
      incurred by, or claims asserted against, such Contributor by reason
      of your accepting any such warranty or additional liability.

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""AMK HTTP client: GET/POST/PUT/DELETE, streaming, AMK_ENV host selection."""

from amk_client.api_config import (
    PATH_PARAM_MODELS,
    ApiEntry,
    api_config,
    get_api_entry,
    path_placeholder_names,
    validate_and_format_path,
)
from amk_client.api_path_params import QueryTaskPathParams
from amk_client.api_request import (
    ApiRequest,
    AsyncApiRequest,
    concat_audio,
    concat_media_segments,
    concat_video,
    enhance_video,
    extract_audio,
    image_to_video,
    mux_audio_video,
    normalize_query_task_response,
    query_task,
    query_task_with_polling,
    trim_audio,
    trim_media_duration,
    trim_video,
    understand_video_content,
)
from amk_client.auth_mode import AmkAuthMode
from amk_client.models import (
    ConcatMediaKind,
    ConcatMediaSegmentsAsyncResult,
    ConcatMediaSegmentsRequest,
    ConcatMediaSegmentsResponse,
    EnhanceVideoRequest,
    ExtractAudioRequest,
    ImageToVideoImageItem,
    ImageToVideoRequest,
    MuxAudioVideoRequest,
    QueryTaskNormalizedResult,
    TaskStatusStr,
    TransitionId,
    TrimMediaDurationAsyncResult,
    TrimMediaDurationRequest,
    TrimMediaDurationResponse,
    TrimMediaKind,
)
from amk_client.client import AmkAsyncClient, AmkClient
from amk_client.config import (
    MissingConfigError,
    build_authorization_header,
    default_ark_api_key,
    default_bearer_token,
    normalize_secret,
    require_amk_api_key,
    require_ark_api_key,
    require_credentials_for_mode,
    resolve_base_url,
    resolve_default_auth_mode,
)

__all__ = [
    "AmkAuthMode",
    "AmkAsyncClient",
    "AmkClient",
    "ApiEntry",
    "AsyncApiRequest",
    "ApiRequest",
    "ConcatMediaKind",
    "ConcatMediaSegmentsAsyncResult",
    "concat_audio",
    "concat_media_segments",
    "concat_video",
    "ConcatMediaSegmentsRequest",
    "ConcatMediaSegmentsResponse",
    "EnhanceVideoRequest",
    "enhance_video",
    "ExtractAudioRequest",
    "ImageToVideoImageItem",
    "ImageToVideoRequest",
    "MuxAudioVideoRequest",
    "mux_audio_video",
    "normalize_query_task_response",
    "query_task",
    "query_task_with_polling",
    "QueryTaskNormalizedResult",
    "TaskStatusStr",
    "MissingConfigError",
    "PATH_PARAM_MODELS",
    "QueryTaskPathParams",
    "TrimMediaDurationAsyncResult",
    "TrimMediaDurationRequest",
    "TrimMediaDurationResponse",
    "trim_audio",
    "trim_media_duration",
    "trim_video",
    "TransitionId",
    "understand_video_content",
    "image_to_video",
    "extract_audio",
    "TrimMediaKind",
    "api_config",
    "build_authorization_header",
    "default_ark_api_key",
    "default_bearer_token",
    "get_api_entry",
    "normalize_secret",
    "path_placeholder_names",
    "require_amk_api_key",
    "require_ark_api_key",
    "require_credentials_for_mode",
    "resolve_base_url",
    "resolve_default_auth_mode",
    "validate_and_format_path",
]

"""AMK HTTP API 注册表：路径、方法、鉴权模式。"""

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import re
from typing import Any, Literal, Mapping

from pydantic import BaseModel, ConfigDict

from amk_client.api_path_params import QueryTaskPathParams
from amk_client.auth_mode import AmkAuthMode

HttpMethod = Literal["GET", "POST", "PUT", "DELETE", "PATCH"]

_PATH_PLACEHOLDERS = re.compile(r"\{(\w+)\}")


class ApiEntry(BaseModel):
    """单条 API 配置（不可变）。"""

    model_config = ConfigDict(frozen=True)

    action: str
    auth_mode: AmkAuthMode
    method: HttpMethod = "POST"


# api_name -> 路径参数模型（无路径占位符的 API 不必注册）
PATH_PARAM_MODELS: dict[str, type[BaseModel]] = {
    "query_task": QueryTaskPathParams,
}


api_config: dict[str, ApiEntry] = {
    "understand_video_content": ApiEntry(
        action="/api/v1/chat/completions",
        auth_mode=AmkAuthMode.ARK_AND_MEDIAKIT,
        method="POST",
    ),
    "trim_audio": ApiEntry(
        action="/api/v1/tools/trim-audio",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "trim_video": ApiEntry(
        action="/api/v1/tools/trim-video",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    # 调用节点聚合：统一由 concat_media_segments(type=audio|video) 分发到底层 concat_* API
    "concat_audio": ApiEntry(
        action="/api/v1/tools/concat-audio",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "concat_video": ApiEntry(
        action="/api/v1/tools/concat-video",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "mux_audio_video": ApiEntry(
        action="/api/v1/tools/mux-audio-video",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "image_to_video": ApiEntry(
        action="/api/v1/tools/image-to-video",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "extract_audio": ApiEntry(
        action="/api/v1/tools/extract-audio",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "enhance_video": ApiEntry(
        action="/api/v1/tools/enhance-video",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="POST",
    ),
    "query_task": ApiEntry(
        action="/api/v1/tasks/{task_id}",
        auth_mode=AmkAuthMode.MEDIAKIT_ONLY,
        method="GET",
    ),
}


def get_api_entry(name: str) -> ApiEntry:
    """按注册名取配置；未知名称抛 ``KeyError``。"""
    return api_config[name]


def path_placeholder_names(action: str) -> tuple[str, ...]:
    return tuple(_PATH_PLACEHOLDERS.findall(action))


def validate_and_format_path(
    api_name: str,
    entry: ApiEntry,
    path_params: Mapping[str, Any] | BaseModel | None,
) -> str:
    """用 Pydantic 校验路径参数（若已注册模型），再 ``str.format`` 生成最终 path。"""
    names = path_placeholder_names(entry.action)
    if not names:
        if path_params:
            raise ValueError(
                f"API {api_name!r} 的 action 无路径占位符，不应传入 path_params: {path_params!r}"
            )
        return entry.action

    raw: Mapping[str, Any]
    if path_params is None:
        raw = {}
    elif isinstance(path_params, BaseModel):
        raw = path_params.model_dump()
    else:
        raw = path_params

    model_cls = PATH_PARAM_MODELS.get(api_name)
    if model_cls is not None:
        data = model_cls.model_validate(raw).model_dump()
    else:
        data = dict(raw)

    missing = [n for n in names if n not in data]
    if missing:
        raise ValueError(
            f"API {api_name!r} 缺少路径参数: {missing}；已注册模型时可自动校验类型与约束"
        )

    return entry.action.format(**{n: str(data[n]) for n in names})

"""各 API 路径占位符对应的 Pydantic 校验模型（与 ``api_config`` 中 ``action`` 的 ``{name}`` 一致）。"""

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from pydantic import BaseModel, ConfigDict, Field


class QueryTaskPathParams(BaseModel):
    """``GET /api/v1/tasks/{task_id}``"""

    model_config = ConfigDict(str_strip_whitespace=True)

    task_id: str = Field(..., min_length=1, description="任务 ID")

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""按 ``api_config`` 发起请求；路径参数经 Pydantic 校验。"""

from __future__ import annotations

import asyncio
import json
import os
import uuid
from typing import Any, Mapping

import httpx
from pydantic import BaseModel

from amk_client.api_config import get_api_entry, validate_and_format_path
from amk_client.client import AmkAsyncClient, AmkClient
from amk_client.config import MissingConfigError, normalize_secret
from amk_client.models import (
    ConcatMediaSegmentsAsyncResult,
    ConcatMediaSegmentsRequest,
    ConcatMediaSegmentsResponse,
    EnhanceVideoRequest,
    ExtractAudioRequest,
    ImageToVideoRequest,
    MuxAudioVideoRequest,
    QueryTaskNormalizedResult,
    TaskStatusStr,
    TrimMediaDurationAsyncResult,
    TrimMediaDurationRequest,
    TrimMediaDurationResponse,
)


class ApiRequest:
    """封装单个注册 API：路径参数校验 + 按配置的 method/auth_mode 调用 ``AmkClient``。"""

    def __init__(self, api_name: str, *, client: AmkClient | None = None) -> None:
        self.api_name = api_name
        self.entry = get_api_entry(api_name)
        self.client = client if client is not None else AmkClient()

    def build_path(
        self,
        path_params: Mapping[str, Any] | BaseModel | None = None,
    ) -> str:
        return validate_and_format_path(self.api_name, self.entry, path_params)

    def request(
        self,
        *,
        path_params: Mapping[str, Any] | BaseModel | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        path = self.build_path(path_params)
        method = self.entry.method.upper()
        auth_mode = self.entry.auth_mode
        force_new_client_token = bool(kwargs.pop("force_new_client_token", False))
        kwargs = _inject_client_token_if_needed(
            self.api_name, kwargs, force_new_client_token=force_new_client_token
        )
        if method in ("GET", "POST", "PUT", "DELETE"):
            return getattr(self.client, method.lower())(
                path, auth_mode=auth_mode, **kwargs
            )
        return self.client.request(method, path, auth_mode=auth_mode, **kwargs)


class AsyncApiRequest:
    """异步版 API 调用封装。"""

    def __init__(self, api_name: str, *, client: AmkAsyncClient | None = None) -> None:
        self.api_name = api_name
        self.entry = get_api_entry(api_name)
        self.client = client if client is not None else AmkAsyncClient()

    def build_path(
        self,
        path_params: Mapping[str, Any] | BaseModel | None = None,
    ) -> str:
        return validate_and_format_path(self.api_name, self.entry, path_params)

    async def request(
        self,
        *,
        path_params: Mapping[str, Any] | BaseModel | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        path = self.build_path(path_params)
        method = self.entry.method.upper()
        auth_mode = self.entry.auth_mode
        force_new_client_token = bool(kwargs.pop("force_new_client_token", False))
        kwargs = _inject_client_token_if_needed(
            self.api_name, kwargs, force_new_client_token=force_new_client_token
        )
        if method in ("GET", "POST", "PUT", "DELETE"):
            return await getattr(self.client, method.lower())(
                path, auth_mode=auth_mode, **kwargs
            )
        return await self.client.request(method, path, auth_mode=auth_mode, **kwargs)


def _inject_client_token_if_needed(
    api_name: str,
    kwargs: Mapping[str, Any],
    *,
    force_new_client_token: bool = False,
) -> dict[str, Any]:
    """为非视频理解请求按开关注入幂等 token。"""
    patched = dict(kwargs)
    if api_name == "understand_video_content":
        return patched
    enabled = _is_client_token_enabled()
    if not enabled and not force_new_client_token:
        return patched
    payload = patched.get("json")
    if not isinstance(payload, Mapping):
        return patched
    body = dict(payload)
    if force_new_client_token or "client_token" not in body:
        body["client_token"] = uuid.uuid4().hex[:8]
    patched["json"] = body
    return patched


def _is_client_token_enabled() -> bool:
    return (os.getenv("AMK_ENABLE_CLIENT_TOKEN") or "false").strip().lower() == "true"


async def trim_audio(
    payload: TrimMediaDurationRequest,
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = {"audio_url": payload.source, "start_time": payload.start_time}
    if payload.end_time is not None:
        body["end_time"] = payload.end_time
    return await AsyncApiRequest("trim_audio", client=use_client).request(json=body)


async def trim_video(
    payload: TrimMediaDurationRequest,
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = {"video_url": payload.source, "start_time": payload.start_time}
    if payload.end_time is not None:
        body["end_time"] = payload.end_time
    return await AsyncApiRequest("trim_video", client=use_client).request(json=body)


async def trim_media_duration(
    payload: TrimMediaDurationRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> TrimMediaDurationAsyncResult:
    """聚合方法：按 type 分发 trim_audio/trim_video，异常时返回 error。"""
    try:
        body = (
            payload
            if isinstance(payload, TrimMediaDurationRequest)
            else TrimMediaDurationRequest.model_validate(payload)
        )
        if body.type == "audio":
            res = await trim_audio(body, client=client)
        else:
            res = await trim_video(body, client=client)
        res.raise_for_status()
        data = TrimMediaDurationResponse.model_validate(res.json())
        return TrimMediaDurationAsyncResult(
            task_id=data.task_id, request_id=data.request_id
        )
    except Exception as exc:  # noqa: BLE001 - 统一返回 error 给调用方
        return TrimMediaDurationAsyncResult(error=str(exc))


async def concat_audio(
    payload: ConcatMediaSegmentsRequest,
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = {"audio_urls": payload.sources}
    return await AsyncApiRequest("concat_audio", client=use_client).request(json=body)


async def concat_video(
    payload: ConcatMediaSegmentsRequest,
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = {"video_urls": payload.sources}
    if payload.transitions is not None:
        body["transitions"] = payload.transitions
    return await AsyncApiRequest("concat_video", client=use_client).request(json=body)


async def concat_media_segments(
    payload: ConcatMediaSegmentsRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> ConcatMediaSegmentsAsyncResult:
    """聚合方法：按 type 分发 concat_audio/concat_video，异常时返回 error。"""
    try:
        body = (
            payload
            if isinstance(payload, ConcatMediaSegmentsRequest)
            else ConcatMediaSegmentsRequest.model_validate(payload)
        )
        if body.type == "audio":
            res = await concat_audio(body, client=client)
        else:
            res = await concat_video(body, client=client)
        res.raise_for_status()
        data = ConcatMediaSegmentsResponse.model_validate(res.json())
        return ConcatMediaSegmentsAsyncResult(
            task_id=data.task_id, request_id=data.request_id
        )
    except Exception as exc:  # noqa: BLE001 - 统一返回 error 给调用方
        return ConcatMediaSegmentsAsyncResult(error=str(exc))


async def understand_video_content(
    prompt: str,
    video_url: str,
    fps: int | float,
    max_frames: int | None = None,
    *,
    client: AmkAsyncClient | None = None,
) -> list[dict[str, Any]]:
    """
    视频理解（固定入参版本）：
    - 接收 prompt、video_url、fps，支持可选 max_frames
    - 内部组装 chat.completions payload
    - 仅返回 choices[].message
    """
    if not prompt.strip():
        raise ValueError("prompt 不能为空")
    if not (video_url.startswith("http://") or video_url.startswith("https://")):
        raise ValueError("video_url 须为 http:// 或 https:// 开头")
    if fps <= 0:
        raise ValueError("fps 必须大于 0")
    if max_frames is not None and max_frames <= 0:
        raise ValueError("max_frames 必须大于 0")
    model_id = normalize_secret(os.environ.get("ARK_MODEL_ID"))
    if model_id is None:
        raise MissingConfigError.missing_env(
            "ARK_MODEL_ID",
            purpose="视频理解 chat.completions 的 model 字段。",
            value_placeholder="<Ark Model ID>",
        )

    use_client = client if client is not None else AmkAsyncClient()
    video_payload: dict[str, Any] = {
        "url": video_url,
        "fps": fps,
    }
    if max_frames is not None:
        video_payload["max_frames"] = max_frames

    request_payload = {
        "model": model_id,
        "messages": [
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": prompt},
                    {
                        "type": "video_url",
                        "video_url": video_payload,
                    },
                ],
            }
        ],
        "stream": False,
    }
    resp = await AsyncApiRequest("understand_video_content", client=use_client).request(
        json=request_payload
    )
    resp.raise_for_status()
    body = resp.json()
    data = body.get("data") if isinstance(body, Mapping) else None
    container = data if isinstance(data, Mapping) else body
    choices = container.get("choices") if isinstance(container, Mapping) else None
    if not isinstance(choices, list):
        raise ValueError("understand_video_content 响应缺少 choices")

    messages: list[dict[str, Any]] = []
    for choice in choices:
        if isinstance(choice, Mapping) and isinstance(choice.get("message"), Mapping):
            messages.append(dict(choice["message"]))
    return messages


async def mux_audio_video(
    payload: MuxAudioVideoRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = (
        payload.model_dump(mode="json", exclude_none=True)
        if isinstance(payload, MuxAudioVideoRequest)
        else MuxAudioVideoRequest.model_validate(payload).model_dump(
            mode="json", exclude_none=True
        )
    )
    return await AsyncApiRequest("mux_audio_video", client=use_client).request(
        json=body
    )


async def image_to_video(
    payload: ImageToVideoRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = (
        payload.model_dump(mode="json", exclude_none=True)
        if isinstance(payload, ImageToVideoRequest)
        else ImageToVideoRequest.model_validate(payload).model_dump(
            mode="json", exclude_none=True
        )
    )
    return await AsyncApiRequest("image_to_video", client=use_client).request(json=body)


async def extract_audio(
    payload: ExtractAudioRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = (
        payload.model_dump(mode="json", exclude_none=True)
        if isinstance(payload, ExtractAudioRequest)
        else ExtractAudioRequest.model_validate(payload).model_dump(
            mode="json", exclude_none=True
        )
    )
    return await AsyncApiRequest("extract_audio", client=use_client).request(json=body)


async def query_task(
    task_id: str,
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    return await AsyncApiRequest("query_task", client=use_client).request(
        path_params={"task_id": task_id}
    )


def _extract_play_url(
    task_type: str | None, parsed_result: Mapping[str, Any]
) -> str | None:
    # 优先查找 video_url 和 audio_url，不管任务类型
    video_url = parsed_result.get("video_url") or parsed_result.get("videoa_url")
    audio_url = parsed_result.get("audio_url")
    if video_url:
        return video_url
    if audio_url:
        return audio_url
    return parsed_result.get("play_url")


def normalize_query_task_response(
    payload: Mapping[str, Any],
) -> QueryTaskNormalizedResult:
    # 新接口：status / task_id / result 等在根上；旧接口：包在 Result 里
    if isinstance(payload.get("status"), str):
        task_block: Mapping[str, Any] = payload
    else:
        wrapped = payload.get("Result")
        if not isinstance(wrapped, Mapping):
            raise ValueError("query_task 响应缺少 status 或 Result 对象")
        task_block = wrapped

    status = task_block.get("status")
    if not isinstance(status, str):
        raise ValueError("query_task 响应缺少 status 字段")

    # 新接口通常是根字段；旧接口可能在 Result 内
    task_id = None
    if isinstance(task_block.get("task_id"), str):
        task_id = task_block["task_id"]
    elif isinstance(task_block.get("TaskId"), str):
        task_id = task_block["TaskId"]
    elif isinstance(payload.get("task_id"), str):
        task_id = payload["task_id"]
    elif isinstance(payload.get("TaskId"), str):
        task_id = payload["TaskId"]

    task_type = (
        task_block.get("task_type")
        if isinstance(task_block.get("task_type"), str)
        else None
    )
    request_id = None
    if isinstance(task_block.get("request_id"), str):
        request_id = task_block["request_id"]
    else:
        meta = payload.get("ResponseMetadata")
        if isinstance(meta, Mapping) and isinstance(meta.get("RequestId"), str):
            request_id = meta["RequestId"]

    error = task_block.get("error")
    if status in ("failed", "canceled"):
        raise RuntimeError(str(error or f"task {status}"))

    result_raw = task_block.get("result")
    parsed_result: dict[str, Any] = {}
    if isinstance(result_raw, str) and result_raw.strip():
        parsed = json.loads(result_raw)
        if isinstance(parsed, Mapping):
            parsed_result = dict(parsed)
    elif isinstance(result_raw, Mapping):
        parsed_result = dict(result_raw)

    return QueryTaskNormalizedResult(
        task_id=task_id,
        duration=parsed_result.get("duration"),
        play_url=_extract_play_url(task_type, parsed_result),
        request_id=request_id,
        status=status,  # type: ignore[arg-type]
        task_type=task_type,
    )


async def query_task_with_polling(
    task_id: str,
    interval: int = 5,
    max_retries: int = 6,
    *,
    client: AmkAsyncClient | None = None,
) -> QueryTaskNormalizedResult:
    """
    轮询任务直到完成或失败。

    - status=running/queued: 继续轮询
    - status=completed: 返回标准化结果
    - status=failed/canceled: 抛异常
    """
    if interval <= 0:
        raise ValueError("interval 必须大于 0")
    if max_retries <= 0:
        raise ValueError("max_retries 必须大于 0")

    for attempt in range(max_retries):
        resp = await query_task(task_id, client=client)
        resp.raise_for_status()
        normalized = normalize_query_task_response(resp.json())
        status: TaskStatusStr = normalized.status
        if status == "completed":
            return normalized
        if status in ("running", "queued"):
            if attempt < max_retries - 1:
                await asyncio.sleep(interval)
                continue
            raise TimeoutError("task polling timeout")
        raise RuntimeError(f"unsupported task status: {status}")

    raise TimeoutError("task polling timeout")


async def enhance_video(
    payload: EnhanceVideoRequest | Mapping[str, Any],
    *,
    client: AmkAsyncClient | None = None,
) -> httpx.Response:
    use_client = client if client is not None else AmkAsyncClient()
    body = (
        payload.model_dump(mode="json", exclude_none=True)
        if isinstance(payload, EnhanceVideoRequest)
        else EnhanceVideoRequest.model_validate(payload).model_dump(
            mode="json", exclude_none=True
        )
    )
    return await AsyncApiRequest("enhance_video", client=use_client).request(json=body)

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from enum import Enum


class AmkAuthMode(str, Enum):
    """AMK HTTP API 鉴权注册方式（与网关约定一致）。"""

    #: ``Authorization: Bearer <MediaKit API Key>``
    MEDIAKIT_ONLY = "mediakit_only"
    #: ``Authorization: Bearer <火山方舟 API Key>/<MediaKit API Key>``
    ARK_AND_MEDIAKIT = "ark_and_mediakit"

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from collections.abc import Iterator, Mapping
from contextlib import contextmanager
from typing import Any

import httpx

from amk_client.auth_mode import AmkAuthMode
from amk_client.config import (
    build_authorization_header,
    default_ark_api_key,
    default_bearer_token,
    normalize_secret,
    require_credentials_for_mode,
    resolve_base_url,
    resolve_default_auth_mode,
)


def _load_dotenv() -> None:
    try:
        from dotenv import find_dotenv, load_dotenv

        path = find_dotenv(usecwd=True)
        if path:
            load_dotenv(path)
    except ImportError:
        pass


class AmkClient:
    """
    Sync HTTP client for AMK APIs.

    **鉴权模式按「单次请求」指定**（与网关上各 API 注册方式一致），在
    ``request`` / ``get`` / ``post`` / ``put`` / ``delete`` / ``stream`` /
    ``iter_*`` 上传入 ``auth_mode=...``。

    - ``AmkAuthMode.MEDIAKIT_ONLY``：``Authorization: Bearer <MediaKit API Key>``
    - ``AmkAuthMode.ARK_AND_MEDIAKIT``：
      ``Authorization: Bearer <火山方舟 API Key>/<MediaKit API Key>``

    若某次调用**未传** ``auth_mode``，则使用构造参数 ``auth_mode``（默认模式）；
    构造时也未指定时，再读环境变量 ``AMK_AUTH_MODE``。

    密钥来自 ``AMK_API_KEY`` / ``ARK_API_KEY`` 或构造参数 ``api_key`` / ``ark_api_key``。

    - Base URL 默认由 ``AMK_ENV`` 与当次 ``auth_mode`` 共同决定（同一环境下不同鉴权模式可命中不同 host 组）；也可用构造参数 ``base_url`` 覆盖。
    - ``require_api_key=True``（默认）时，**每次请求**按当次 ``auth_mode`` 校验所需密钥。
    - ``require_api_key=False`` 时不校验；有密钥则仍按当次模式组 ``Authorization``。
    - 单次请求可传 ``headers`` 覆盖 ``Authorization``。
    """

    def __init__(
        self,
        *,
        auth_mode: AmkAuthMode | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        ark_api_key: str | None = None,
        require_api_key: bool = True,
        timeout: float | httpx.Timeout | None = 30.0,
        extra_headers: Mapping[str, str] | None = None,
    ) -> None:
        _load_dotenv()
        # 仅当具体请求未传 auth_mode 时使用（非「全局业务配置」，而是默认值）
        self._default_auth_mode = (
            auth_mode if auth_mode is not None else resolve_default_auth_mode()
        )
        self._base_override = base_url.rstrip("/") if base_url else None
        self._mediakit_explicit = api_key
        self._ark_explicit = ark_api_key
        self._require_api_key = require_api_key
        self._timeout = timeout
        self._extra_headers = dict(extra_headers) if extra_headers else {}

    def _mediakit_resolved(self) -> str | None:
        if self._mediakit_explicit is not None:
            return normalize_secret(self._mediakit_explicit)
        return default_bearer_token()

    def _ark_resolved(self) -> str | None:
        if self._ark_explicit is not None:
            return normalize_secret(self._ark_explicit)
        return default_ark_api_key()

    def _authorization_for_mode(self, mode: AmkAuthMode) -> str | None:
        mk = self._mediakit_resolved()
        if mode == AmkAuthMode.MEDIAKIT_ONLY:
            if not mk:
                return None
            return build_authorization_header(
                mediakit_api_key=mk, mode=AmkAuthMode.MEDIAKIT_ONLY
            )
        ark = self._ark_resolved()
        if not mk or not ark:
            return None
        return build_authorization_header(
            mediakit_api_key=mk,
            ark_api_key=ark,
            mode=AmkAuthMode.ARK_AND_MEDIAKIT,
        )

    def _effective_auth_mode(self, auth_mode: AmkAuthMode | None) -> AmkAuthMode:
        return auth_mode if auth_mode is not None else self._default_auth_mode

    def _merge_headers(
        self,
        call_headers: Mapping[str, str] | None,
        *,
        auth_mode: AmkAuthMode | None = None,
    ) -> dict[str, str]:
        mode = self._effective_auth_mode(auth_mode)
        if self._require_api_key:
            require_credentials_for_mode(
                mode,
                mediakit_from_constructor=self._mediakit_explicit,
                ark_from_constructor=self._ark_explicit,
            )
        merged: dict[str, str] = {**self._extra_headers}
        auth = self._authorization_for_mode(mode)
        if auth:
            merged.setdefault("Authorization", auth)
        # 默认使用 JSON 请求体；调用方可在单次请求 headers 中覆盖
        merged.setdefault("Content-Type", "application/json")
        if call_headers:
            merged.update(call_headers)
        return merged

    def _url(self, path: str, *, auth_mode: AmkAuthMode | None = None) -> str:
        p = path if path.startswith("/") else f"/{path}"
        mode = self._effective_auth_mode(auth_mode)
        base = self._base_override or resolve_base_url(
            auth_mode=mode, fallback_env="prod"
        )
        return f"{base}{p}"

    def _client(self) -> httpx.Client:
        return httpx.Client(timeout=self._timeout)

    def request(
        self,
        method: str,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        with self._client() as client:
            return client.request(
                method.upper(),
                self._url(path, auth_mode=auth_mode),
                headers=self._merge_headers(headers, auth_mode=auth_mode),
                **kwargs,
            )

    def get(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request("GET", path, auth_mode=auth_mode, headers=headers, **kwargs)

    def post(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            "POST", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

    def put(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request("PUT", path, auth_mode=auth_mode, headers=headers, **kwargs)

    def delete(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return self.request(
            "DELETE", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

    @contextmanager
    def stream(
        self,
        method: str,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> Iterator[httpx.Response]:
        """Stream response body (SSE/chunked). Caller should read ``iter_bytes`` / ``iter_lines``."""
        with self._client() as client:
            with client.stream(
                method.upper(),
                self._url(path, auth_mode=auth_mode),
                headers=self._merge_headers(headers, auth_mode=auth_mode),
                **kwargs,
            ) as response:
                yield response

    def iter_bytes(
        self,
        method: str,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        chunk_size: int | None = None,
        **kwargs: Any,
    ) -> Iterator[bytes]:
        """Convenience: yield byte chunks from a streaming request."""
        with self.stream(
            method, path, auth_mode=auth_mode, headers=headers, **kwargs
        ) as resp:
            yield from resp.iter_bytes(chunk_size=chunk_size)

    def iter_lines(
        self,
        method: str,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> Iterator[str]:
        """Convenience: yield text lines (e.g. SSE) from a streaming request."""
        with self.stream(
            method, path, auth_mode=auth_mode, headers=headers, **kwargs
        ) as resp:
            yield from resp.iter_lines()


class AmkAsyncClient(AmkClient):
    """Async HTTP client for AMK APIs."""

    def _async_client(self) -> httpx.AsyncClient:
        return httpx.AsyncClient(timeout=self._timeout)

    async def request(
        self,
        method: str,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        async with self._async_client() as client:
            return await client.request(
                method.upper(),
                self._url(path, auth_mode=auth_mode),
                headers=self._merge_headers(headers, auth_mode=auth_mode),
                **kwargs,
            )

    async def get(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            "GET", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

    async def post(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            "POST", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

    async def put(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            "PUT", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

    async def delete(
        self,
        path: str,
        *,
        auth_mode: AmkAuthMode | None = None,
        headers: Mapping[str, str] | None = None,
        **kwargs: Any,
    ) -> httpx.Response:
        return await self.request(
            "DELETE", path, auth_mode=auth_mode, headers=headers, **kwargs
        )

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import os

from amk_client.auth_mode import AmkAuthMode

HOSTS_BY_AUTH_MODE: dict[AmkAuthMode, dict[str, str]] = {
    AmkAuthMode.MEDIAKIT_ONLY: {
        "prod": "https://amk.cn-beijing.volces.com",
    },
    AmkAuthMode.ARK_AND_MEDIAKIT: {
        "prod": "https://amk-ark.cn-beijing.volces.com",
    },
}

DEFAULT_AMK_ENV = "prod"
DEFAULT_AMK_AUTH_MODE = AmkAuthMode.MEDIAKIT_ONLY

# 控制台入口，供缺失/错误配置时的 reference_url
AMK_CONSOLE_REFERENCE = "https://console.volcengine.com/imp/ai-mediakit/"
ARK_CONSOLE_REFERENCE = "https://console.volcengine.com/ark/region:ark+cn-beijing/model/detail?Id=doubao-seed-1-8"


class MissingConfigError(ValueError):
    """缺少或非法配置。``var`` 供程序判断；``args[0]`` 为结构化说明，便于人与 Agent 阅读。"""

    var: str

    @classmethod
    def missing_env(
        cls,
        name: str,
        *,
        purpose: str,
        reference_url: str | None = None,
        value_placeholder: str = "<值>",
    ) -> MissingConfigError:
        steps = (
            f"在仓库根目录的 `.env` 增加一行: {name}={value_placeholder}",
            f"或在本终端执行: export {name}={value_placeholder}（再重新运行命令）",
            "若由 Agent 处理：向用户索取真实密钥后，直接编辑仓库根目录 `.env` 写入该键值，"
            f"或在执行命令的 shell 中 `export {name}=<密钥>`，然后重试；不要仅复述本错误而不改配置。",
        )
        lines = [
            f"missing_env_var: {name}",
            f"purpose: {purpose}",
            "fix:",
            *[f"  - {s}" for s in steps],
            "agent_instruction: Obtain the secret from the user, then add `"
            f"{name}=<actual_value>` to repo-root `.env` or export it in the shell, and retry; do not stop after parroting this error.",
        ]
        if reference_url:
            lines.append(f"reference_url: {reference_url}")
        msg = "\n".join(lines)
        err = cls(msg)
        err.var = name
        return err

    @classmethod
    def invalid_env_value(
        cls,
        name: str,
        value: str,
        *,
        allowed: tuple[str, ...],
        purpose: str,
    ) -> MissingConfigError:
        lines = [
            f"invalid_env_var: {name}",
            f"given_value: {value!r}",
            f"allowed: {', '.join(allowed)}",
            f"purpose: {purpose}",
            "fix:",
            "  - 将环境变量设为 allowed 中之一；若未设置该变量，将使用默认 `prod`。",
            f"agent_instruction: Edit repo-root `.env` (or export) so `{name}` is one of: "
            f"{', '.join(allowed)}; or remove `{name}` to use default `prod`, then retry.",
        ]
        msg = "\n".join(lines)
        err = cls(msg)
        err.var = name
        return err


def _read_amk_env(*, auth_mode: AmkAuthMode) -> str:
    raw = os.environ.get("AMK_ENV")
    if raw is None or not raw.strip():
        return DEFAULT_AMK_ENV
    v = raw.strip().lower()
    allowed_envs = HOSTS_BY_AUTH_MODE[auth_mode]
    if v not in allowed_envs:
        raise MissingConfigError.invalid_env_value(
            "AMK_ENV",
            v,
            allowed=tuple(sorted(allowed_envs)),
            purpose="选择 AMK 网关环境：当前仅支持 `prod`（线上环境）。",
        )
    return v


def resolve_default_auth_mode() -> AmkAuthMode:
    """从 ``AMK_AUTH_MODE`` 解析 **AmkClient 的默认** 鉴权模式（仅当单次请求未传 ``auth_mode``、且构造时也未指定 ``auth_mode`` 时生效）；未设置环境变量时为 ``mediakit_only``。

    合法取值：``mediakit_only``（或 ``mediakit`` / ``single``）、
    ``ark_and_mediakit``（或 ``ark_mediakit`` / ``dual`` / ``combined``）。
    """
    raw = (os.environ.get("AMK_AUTH_MODE") or "").strip().lower()
    if raw in ("", "mediakit_only", "mediakit", "single"):
        return AmkAuthMode.MEDIAKIT_ONLY
    if raw in ("ark_and_mediakit", "ark_mediakit", "dual", "combined", "ark+mediakit"):
        return AmkAuthMode.ARK_AND_MEDIAKIT
    raise MissingConfigError.invalid_env_value(
        "AMK_AUTH_MODE",
        raw or "(empty)",
        allowed=("mediakit_only", "ark_and_mediakit"),
        purpose="注册 API 时选择的鉴权方式：仅 MediaKit Key，或 方舟 Key + MediaKit Key。",
    )


def build_authorization_header(
    *,
    mediakit_api_key: str,
    ark_api_key: str | None = None,
    mode: AmkAuthMode = DEFAULT_AMK_AUTH_MODE,
) -> str:
    """组装 ``Authorization`` 请求头的完整值（含 ``Bearer `` 前缀）。"""
    if mode == AmkAuthMode.MEDIAKIT_ONLY:
        return f"Bearer {mediakit_api_key}"
    if ark_api_key is None:
        raise ValueError("ark_and_mediakit 模式下必须提供 ark_api_key")
    return f"Bearer {ark_api_key}/{mediakit_api_key}"


def resolve_base_url(
    *,
    auth_mode: AmkAuthMode = DEFAULT_AMK_AUTH_MODE,
    fallback_env: str = DEFAULT_AMK_ENV,
) -> str:
    """Return API base URL for current ``AMK_ENV`` and auth mode.

    若显式设置了非法的 ``AMK_ENV``，抛出 ``MissingConfigError``。
    """
    env = _read_amk_env(auth_mode=auth_mode)
    hosts = HOSTS_BY_AUTH_MODE[auth_mode]
    return hosts.get(env, hosts[fallback_env])


def normalize_secret(value: str | None) -> str | None:
    """非空、非 ``null`` 的字符串视为有效密钥；否则 ``None``。"""
    if value is None:
        return None
    stripped = value.strip()
    if not stripped or stripped.lower() == "null":
        return None
    return stripped


def default_bearer_token() -> str | None:
    """Token from ``AMK_API_KEY`` (empty or ``null`` ignored)."""
    return normalize_secret(os.environ.get("AMK_API_KEY"))


def require_amk_api_key() -> str:
    """返回可用的 MediaKit / AMK ``AMK_API_KEY``；缺失时抛出 ``MissingConfigError``。"""
    key = default_bearer_token()
    if key is None:
        raise MissingConfigError.missing_env(
            "AMK_API_KEY",
            purpose="调用 AMK（MediaKit）HTTP API 时的 Bearer 鉴权。",
            reference_url=AMK_CONSOLE_REFERENCE,
            value_placeholder="<MediaKit API Key>",
        )
    return key


def default_ark_api_key() -> str | None:
    """方舟 ``ARK_API_KEY``；未设置、空或 ``null`` 时返回 ``None``。"""
    return normalize_secret(os.environ.get("ARK_API_KEY"))


def require_ark_api_key() -> str:
    """返回可用的方舟密钥；缺失时抛出 ``MissingConfigError``。"""
    key = default_ark_api_key()
    if key is None:
        raise MissingConfigError.missing_env(
            "ARK_API_KEY",
            purpose="方舟 + MediaKit 双 Key 鉴权时，作为 ``Authorization: Bearer <方舟>/<MediaKit>`` 的前段。",
            reference_url=ARK_CONSOLE_REFERENCE,
            value_placeholder="<Ark API Key>",
        )
    return key


def require_credentials_for_mode(
    mode: AmkAuthMode,
    *,
    mediakit_from_constructor: str | None,
    ark_from_constructor: str | None,
) -> None:
    """按鉴权模式校验必填密钥（构造参数或环境变量）；缺一则抛 ``MissingConfigError``。"""
    mk = (
        normalize_secret(mediakit_from_constructor)
        if mediakit_from_constructor is not None
        else default_bearer_token()
    )
    if mode == AmkAuthMode.MEDIAKIT_ONLY:
        if (
            mediakit_from_constructor is not None
            and normalize_secret(mediakit_from_constructor) is None
        ):
            raise MissingConfigError.missing_env(
                "AMK_API_KEY",
                purpose="单 MediaKit Key 鉴权，或双 Key 模式下的 MediaKit 段。",
                reference_url=AMK_CONSOLE_REFERENCE,
                value_placeholder="<MediaKit API Key>",
            )
        if mk is None:
            require_amk_api_key()
        return

    # ARK_AND_MEDIAKIT
    if (
        mediakit_from_constructor is not None
        and normalize_secret(mediakit_from_constructor) is None
    ):
        raise MissingConfigError.missing_env(
            "AMK_API_KEY",
            purpose="双 Key 鉴权：``Authorization: Bearer <方舟 API Key>/<MediaKit API Key>`` 的后段。",
            reference_url=AMK_CONSOLE_REFERENCE,
            value_placeholder="<MediaKit API Key>",
        )
    if mk is None:
        require_amk_api_key()

    ark = (
        normalize_secret(ark_from_constructor)
        if ark_from_constructor is not None
        else default_ark_api_key()
    )
    if (
        ark_from_constructor is not None
        and normalize_secret(ark_from_constructor) is None
    ):
        raise MissingConfigError.missing_env(
            "ARK_API_KEY",
            purpose="双 Key 鉴权：``Authorization: Bearer <方舟 API Key>/<MediaKit API Key>`` 的前段。",
            reference_url=ARK_CONSOLE_REFERENCE,
            value_placeholder="<Ark API Key>",
        )
    if ark is None:
        require_ark_api_key()

"""AMK 相关 Pydantic 模型（请求体、响应体、路径参数等，随 API 扩展）。"""

#!/usr/bin/env python3
# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

from typing import Any, Literal

from pydantic import (
    AliasChoices,
    BaseModel,
    ConfigDict,
    Field,
    ValidationInfo,
    field_validator,
)

TrimMediaKind = Literal["audio", "video"]


class TrimMediaDurationRequest(BaseModel):
    """``POST /api/v1/tools/trim_media_duration`` 请求体。

    Args:
        type (str): **必选字段**，裁剪类型。``audio`` | ``video``。
        source (str): **必选字段**，待剪切的资源 URL，支持 ``http://`` 或 ``https://`` 格式。
        start_time (int | float): **非必选字段**，裁剪开始时间，默认为 ``0``，表示从头开始裁剪；
            最多保留 2 位小数，单位：秒。
        end_time (int | float | None): **非必选字段**，裁剪结束时间，默认为片源结尾；
            最多保留 2 位小数，单位：秒；省略或 ``null`` 表示至片尾。
    """

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    type: TrimMediaKind = Field(
        ...,
        description="裁剪类型（与调用侧聚合节点一致）：audio | video",
    )
    source: str = Field(..., min_length=1, description="待裁剪资源 URL")
    start_time: float = Field(
        default=0.0, ge=0, description="裁剪开始时间（秒），默认从头"
    )
    end_time: float | None = Field(
        default=None,
        ge=0,
        description="裁剪结束时间（秒），省略表示至片源结尾",
    )

    @field_validator("type", mode="before")
    @classmethod
    def type_must_be_audio_or_video_str(cls, v: Any) -> Any:
        if v is None:
            raise ValueError('type 不能为空；须为字符串 "audio" 或 "video"')
        if not isinstance(v, str):
            raise ValueError(
                f'type 须为字符串类型（取值 "audio" | "video"），当前为 {type(v).__name__}'
            )
        return v

    @field_validator("source", mode="before")
    @classmethod
    def source_must_be_str(cls, v: Any) -> Any:
        if v is None:
            raise ValueError(
                "source 不能为空；须为字符串类型，且为 http:// 或 https:// 开头的资源 URL"
            )
        if not isinstance(v, str):
            raise ValueError(
                f"source 须为字符串类型（资源 URL），当前为 {type(v).__name__}；"
                "若为多段资源请使用对应接口的数组字段，勿把列表误传到 source"
            )
        return v

    @field_validator("source")
    @classmethod
    def source_must_be_http(cls, v: str) -> str:
        s = v.strip()
        if not (s.startswith("http://") or s.startswith("https://")):
            raise ValueError("source 须为 http:// 或 https:// 开头的 URL 字符串")
        return s

    @field_validator("start_time", "end_time", mode="before")
    @classmethod
    def time_fields_must_be_number_or_omit(cls, v: Any, info: ValidationInfo) -> Any:
        field = info.field_name
        if v is None:
            if field == "end_time":
                return None
            raise ValueError(
                "start_time 不能为空（勿传 null）；须为数字类型（秒，≥0），"
                "若从片头开始请省略该字段或传 0"
            )
        if not isinstance(v, (int, float)):
            raise ValueError(f"{field} 须为数字类型（秒），当前为 {type(v).__name__}")
        return v

    @field_validator("start_time", "end_time")
    @classmethod
    def at_most_two_decimal_places(cls, v: float | None) -> float | None:
        if v is None:
            return v
        text = f"{v:.4f}".rstrip("0").rstrip(".")
        if "." in text:
            _, frac = text.split(".", 1)
            if len(frac) > 2:
                raise ValueError("时间最多保留 2 位小数（单位：秒）")
        return v

    def to_api_json(self) -> dict[str, Any]:
        """发给网关的 JSON：省略 ``end_time`` 表示默认到片尾。"""
        return self.model_dump(mode="json", by_alias=True, exclude_none=True)


class TrimMediaDurationResponse(BaseModel):
    """``trim_media_duration`` 返回中的任务信息。

    Returns:
        task_id: 任务查询 id。
        request_id: 日志 id（兼容历史网关字段名 ``requst_id``）。
    """

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    task_id: str = Field(..., description="任务查询 id")
    request_id: str = Field(
        ...,
        validation_alias=AliasChoices("request_id", "requst_id"),
        description="日志 id",
    )

    @field_validator("task_id", "request_id", mode="before")
    @classmethod
    def response_string_ids(cls, v: Any, info: ValidationInfo) -> Any:
        name = info.field_name
        if v is None:
            raise ValueError(f"{name} 不能为空；解析响应时期望非空字符串")
        if not isinstance(v, str):
            raise ValueError(f"{name} 须为字符串类型，当前为 {type(v).__name__}")
        return v


class TrimMediaDurationAsyncResult(BaseModel):
    """调用节点聚合接口的异步返回结构。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    task_id: str | None = Field(default=None, description="任务查询 id")
    request_id: str | None = Field(default=None, description="日志 id")
    error: str | None = Field(default=None, description="报错信息")


ConcatMediaKind = Literal["audio", "video"]
TransitionId = Literal[
    "1182359",
    "1182360",
    "1182358",
    "1182365",
    "1182367",
    "1182368",
    "1182369",
    "1182370",
    "1182373",
    "1182374",
    "1182375",
    "1182378",
]


class ConcatMediaSegmentsRequest(BaseModel):
    """调用层聚合接口 ``concat_media_segments`` 请求体。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    type: ConcatMediaKind = Field(
        ...,
        description="必选，拼接类型：audio | video",
    )
    sources: list[str] = Field(
        ...,
        min_length=1,
        max_length=100,
        description="必选，待拼接资源 URL 列表（http:// 或 https://）",
    )
    transitions: list[TransitionId] | None = Field(
        default=None,
        description="可选，转场效果 ID 列表；音频拼接不支持",
    )

    @field_validator("sources", mode="before")
    @classmethod
    def sources_must_be_string_list(cls, v: Any) -> Any:
        if not isinstance(v, list):
            raise ValueError("sources 须为字符串 URL 列表")
        if not v:
            raise ValueError("sources 不能为空；至少提供 1 个资源 URL")
        return v

    @field_validator("sources")
    @classmethod
    def sources_each_must_be_http_url(cls, values: list[str]) -> list[str]:
        for idx, item in enumerate(values):
            if not isinstance(item, str):
                raise ValueError(f"sources[{idx}] 须为字符串 URL")
            s = item.strip()
            if not (s.startswith("http://") or s.startswith("https://")):
                raise ValueError(f"sources[{idx}] 须为 http:// 或 https:// 开头的 URL")
        return values

    @field_validator("transitions", mode="before")
    @classmethod
    def transitions_must_be_string_list_or_none(cls, v: Any) -> Any:
        if v is None:
            return None
        if not isinstance(v, list):
            raise ValueError("transitions 须为字符串列表或省略")
        for idx, item in enumerate(v):
            if not isinstance(item, str):
                raise ValueError(f"transitions[{idx}] 须为字符串转场 ID")
        return v

    @field_validator("transitions")
    @classmethod
    def audio_should_not_pass_transitions(
        cls, v: list[str] | None, info: ValidationInfo
    ) -> list[str] | None:
        media_type = info.data.get("type")
        if media_type == "audio" and v:
            raise ValueError("音频拼接不支持 transitions，请不要传该字段")
        return v

    def to_api_json(self) -> dict[str, Any]:
        """发给底层 API 的 JSON；不必要字段自动省略。"""
        return self.model_dump(mode="json", by_alias=True, exclude_none=True)


class ConcatMediaSegmentsResponse(BaseModel):
    """``concat_media_segments`` 聚合后的任务返回。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    task_id: str = Field(..., description="任务查询 id")
    request_id: str = Field(
        ...,
        validation_alias=AliasChoices("request_id", "requst_id"),
        description="日志 id",
    )


class ConcatMediaSegmentsAsyncResult(BaseModel):
    """调用层聚合方法的异步返回结构。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    task_id: str | None = Field(default=None, description="任务查询 id")
    request_id: str | None = Field(default=None, description="日志 id")
    error: str | None = Field(default=None, description="报错信息")


class MuxAudioVideoRequest(BaseModel):
    """``mux_audio_video`` 请求体。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    video_url: str = Field(..., description="输入视频 URL")
    audio_url: str = Field(..., description="输入音频 URL")
    is_audio_reserve: bool = Field(default=True, description="是否保留原视频音频")
    is_video_audio_sync: bool = Field(default=False, description="是否对齐音视频时长")
    sync_mode: Literal["video", "audio"] = Field(
        default="video", description="对齐基准"
    )
    sync_method: Literal["speed", "trim"] = Field(
        default="trim", description="对齐方式"
    )


class ImageToVideoImageItem(BaseModel):
    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    image_url: str = Field(..., description="图片 URL")
    duration: float | None = Field(default=3.0, ge=0, description="播放时长（秒）")
    animation_type: (
        Literal[
            "move_up", "move_down", "move_left", "move_right", "zoom_in", "zoom_out"
        ]
        | None
    ) = Field(default=None, description="动画类型")
    animation_in: float | None = Field(
        default=None, ge=0, description="动画开始时间（秒）"
    )
    animation_out: float | None = Field(
        default=None, ge=0, description="动画结束时间（秒）"
    )


class ImageToVideoRequest(BaseModel):
    """``image_to_video`` 请求体。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    images: list[ImageToVideoImageItem] = Field(..., min_length=1, max_length=100)
    transitions: list[TransitionId] | None = Field(
        default=None, description="可选转场效果 ID"
    )


class ExtractAudioRequest(BaseModel):
    """``extract_audio`` 请求体。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    video_url: str = Field(..., description="输入视频 URL")
    format: Literal["mp3", "m4a"] = Field(default="m4a", description="输出音频格式")


class EnhanceVideoRequest(BaseModel):
    """``enhance_video`` 请求体。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    video_url: str = Field(..., description="输入视频 URL")
    tool_version: Literal["standard", "professional"] = Field(
        default="standard",
        description="工具版本：standard 标准版，professional 专业版",
    )
    resolution: (
        Literal["240p", "360p", "480p", "540p", "720p", "1080p", "2k", "4k"] | None
    ) = Field(
        default=None,
        description="目标分辨率；不传/不填写则使用原始分辨率（后端按原分辨率处理）",
    )
    resolution_limit: int | None = Field(
        default=None, ge=64, le=2160, description="目标长宽限制"
    )
    fps: float | None = Field(default=None, gt=0, le=120, description="目标帧率")


TaskStatusStr = Literal["running", "completed", "queued", "failed", "canceled"]


class QueryTaskNormalizedResult(BaseModel):
    """任务查询的标准化返回结构。"""

    model_config = ConfigDict(str_strip_whitespace=True, populate_by_name=True)

    task_id: str | None = Field(default=None, description="任务查询 ID")
    duration: float | None = Field(default=None, description="时长（秒）")
    play_url: str | None = Field(default=None, description="播放地址")
    request_id: str | None = Field(default=None, description="日志 id")
    status: TaskStatusStr = Field(..., description="任务状态")
    task_type: str | None = Field(default=None, description="任务类型")