"""
Hive AI Image Task Model.

Stores output from the AI-Generated & Deepfake Content Detection model.
"""

from typing import Any

from sqlalchemy import (
    TIMESTAMP,
    Column,
    Float,
    ForeignKey,
    Integer,
    String,
    text,
)

from assets.connectors import mysql


class HiveAiImageTask(mysql.AuModel):
    """Table definition for hive_ai_image_task table."""

    __tablename__ = "hive_ai_image_task"

    id = Column(Integer, primary_key=True, autoincrement=True)
    asset_final_id = Column(Integer, ForeignKey("asset_final.id"), nullable=False)
    task_id = Column(String(255), nullable=False)
    class_name = Column(String(255), nullable=False)
    score_value = Column(Float, nullable=False)
    created_at = Column(
        TIMESTAMP, nullable=False, server_default=text("CURRENT_TIMESTAMP")
    )

    def as_dict(self) -> dict[str, Any]:
        """Return object as dict.

        Returns:
            dict: Dictionary representation of object
        """
        return {
            "id": self.id,
            "asset_final_id": self.asset_final_id,
            "task_id": self.task_id,
            "class_name": self.class_name,
            "score_value": self.score_value,
        }


def create_hive_ai_image_task(
    asset_final_id: int,
    task_id: str,
    class_name: str,
    score_value: float,
) -> dict[str, Any]:
    """Create new Hive AI image task entry."""

    image_task = HiveAiImageTask(
        asset_final_id=asset_final_id,
        task_id=task_id,
        class_name=class_name,
        score_value=score_value,  # type: ignore[arg-type]
    )
    with mysql.au_db_session() as session:
        session.add(image_task)
        session.flush()

        return image_task.as_dict()