from __future__ import annotations

from typing import Any

from configs.kling import KlingConfig
from scripts.kling_client import KlingClient, KlingClientError
from scripts.repository import create_task, mark_task_submission_failed, mark_task_submitted
from scripts.validators import validate_payload
from scripts.query_paths import QUERY_PATHS


def _inject_common(payload: dict[str, Any], callback_url: str | None, external_task_id: str | None) -> dict[str, Any]:
    out = dict(payload)
    if callback_url is not None:
        out['callback_url'] = callback_url
    if external_task_id is not None:
        out['external_task_id'] = external_task_id
    return out


def create_kling_task(endpoint_type: str, payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    client = KlingClient(KlingConfig())
    enriched = _inject_common(payload, callback_url, external_task_id)
    validate_payload(endpoint_type, enriched)
    task_id_internal = create_task(scene_id, endpoint_type, enriched, external_task_id)
    try:
        response = client.post(endpoint_type, enriched)
    except KlingClientError as e:
        mark_task_submission_failed(task_id_internal, str(e), e.category, e.status_code, e.response_body)
        raise
    data = response.get('data') or {}
    mark_task_submitted(task_id_internal, data.get('task_id'), response.get('request_id'), response)
    return {'task_id_internal': task_id_internal, 'response': response}


def query_task_status(query_kind: str, url: str | None = None) -> dict[str, Any]:
    client = KlingClient(KlingConfig())
    target = url or QUERY_PATHS.get(query_kind)
    if not target:
        raise ValueError(f'Query path not pinned yet for {query_kind}')
    return client.get_url(target, log_label=f'query_{query_kind}')


def create_text_to_video(payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    return create_kling_task('text2video', payload, callback_url, external_task_id, scene_id)


def create_image_to_video(payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    return create_kling_task('image2video', payload, callback_url, external_task_id, scene_id)


def create_omni_video(payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    return create_kling_task('omni', payload, callback_url, external_task_id, scene_id)


def create_reference_to_video(payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    return create_kling_task('reference2video', payload, callback_url, external_task_id, scene_id)


def create_extend_video(payload: dict[str, Any], callback_url: str | None = None, external_task_id: str | None = None, scene_id: str | None = None) -> dict[str, Any]:
    return create_kling_task('extend', payload, callback_url, external_task_id, scene_id)