"""OWS Assets API client."""

import config
import httpx
from owsclient import OwsClient
from owsclient import M2MTokenManager
from secrets_manager.python_ext import PythonSecretsManager

secrets_manager = PythonSecretsManager(
    application_context=False,
    environment=config.ENVIRONMENT,
    service_name=config.APPLICATION_NAME,
)

client = OwsClient(
    config.ENVIRONMENT,
    config.APPLICATION_NAME,
    retries=config.OWS_RETRIES,
    timeout=httpx.Timeout(config.OWS_TIMEOUT),
    m2m_token_manager=(
        M2MTokenManager(
            secrets_manager=secrets_manager,
            environment=config.TOKEN_MANAGER_ENVIRONMENT,
            service_name=config.APPLICATION_NAME,
        )
    ),
)


def get_presigned_url(asset_final_id: int) -> str:
    """Get a presigned URL for an asset."""
    response = client.get(
        "ows-assets",
        "/v2/asset/download_url",
        params={
            "asset_final_id": asset_final_id,
            "expires_in": 300,
        },
    )
    response.raise_for_status()
    return response.json()["url"]


def post_results(asset_final_id: int, data: dict, overwrite: bool = False):
    """Post results to OWS Assets API."""
    payload = {
        "asset_final_id": asset_final_id,
        "task": _format_raw_hive_task_data(data),
        "segments": _format_raw_hive_segment_data(data),
    }
    if overwrite:
        payload["overwrite"] = True
    response = client.post("ows-assets", "/hive_segments", json=payload)
    response.raise_for_status()


def _format_raw_hive_segment_data(data: dict) -> list[dict]:
    """Format raw hive data to OWS Assets API format."""
    allowed_data = (
        "ai_generated_music",
        "ai_generated_music_vocal",
        "mubert",
        "musicgen",
        "riffusion",
        "stable_audio",
        "suno",
        "udio",
        "yue",
        "minimax",
        "mureka",
        "ace_step",
        "duobao",
        "google",
        "heartmula",
        "loudly",
    )
    # Class names for the same entity might change over time or model to model
    renamed_data = {"stable_audio_open": "stable_audio"}
    return [
        {
            "time": s["time"],
            **{
                renamed_data.get(x["class"], x["class"]): x["score"]
                for x in s["classes"]
                if x["class"] in allowed_data or x["class"] in renamed_data
            },
        }
        for s in data["status"][0]["response"]["output"]
    ]


def _format_raw_hive_task_data(data: dict) -> dict:
    """Format raw hive task data to OWS Assets API format."""
    return {
        "task_id": data["status"][0]["response"]["input"]["id"],
        "model": data["status"][0]["response"]["input"]["model"],
        "model_version": data["status"][0]["response"]["input"]["model_version"],
    }