"""Utils code."""

import os
from typing import Dict
from typing import Iterable

import boto3
from kafka_utils.consumer.deserializer.string import StringDeserializer
from kafka_utils.consumer.deserializer.simple_json import JSONDeserializer
from kafka_utils.consumer.source.mapping import EventSourceMessage
import pydantic
from pydantic import BaseModel


def download_from_s3_url(bucket, key, dir='.'):
    """Download file from s3 to local dir."""
    s3 = boto3.client('s3')
    filename = os.path.basename(key)
    local_path = os.path.join(dir, filename)
    s3.download_file(bucket, key, local_path)
    return local_path


def get_events_from_batch(event: Dict) -> Iterable[Dict]:
    """Flatten records from the source event."""
    string_deserializer = StringDeserializer()
    json_deserializer = JSONDeserializer()
    for _, msk_message in EventSourceMessage(event):
        yield json_deserializer.deserialize(
            string_deserializer.deserialize(msk_message.value))


class FinalAsset(BaseModel):
    """Final_asset model."""

    asset_final_id: int = pydantic.Field(alias='ASSET_FINAL_ID')
    asset_upload_id: int = pydantic.Field(alias='ASSET_UPLOAD_ID')
    asset_type: str = pydantic.Field(alias='ASSET_TYPE')
    bucket: str = pydantic.Field(alias='BUCKET')
    filename: str = pydantic.Field(alias='FILENAME')


class KafkaResultEvent(BaseModel):
    """Final_asset model."""

    final_asset_id: int
    asset_upload_id: int
    bucket: str
    filename: str
    lowlevel_data: dict