# Copyright Contributors to the Amundsen project.
# SPDX-License-Identifier: Apache-2.0

import unittest
from typing import (
    Iterable, Iterator, Union,
)

from amundsen_rds.models import RDSModel
from sqlalchemy import (
    Column, ForeignKey, String,
)
from sqlalchemy.ext.declarative import declarative_base

from databuilder.models.table_serializable import TableSerializable
from databuilder.serializers import mysql_serializer

Base = declarative_base()


class TestTableSerializable(unittest.TestCase):

    def test_table_serializable(self) -> None:
        actors = [Actor('Tom Cruise'), Actor('Meg Ryan')]
        movie = Movie('Top Gun', actors)

        actual = []
        node_row = movie.next_record()
        while node_row:
            actual.append(mysql_serializer.serialize_record(node_row))
            node_row = movie.next_record()

        expected = [
            {
                'rk': 'movie://Top Gun',
                'name': 'Top Gun'
            },
            {
                'rk': 'actor://Tom Cruise',
                'name': 'Tom Cruise'
            },
            {
                'movie_rk': 'movie://Top Gun',
                'actor_rk': 'actor://Tom Cruise'
            },
            {
                'rk': 'actor://Meg Ryan',
                'name': 'Meg Ryan'
            },
            {
                'movie_rk': 'movie://Top Gun',
                'actor_rk': 'actor://Meg Ryan'
            }
        ]

        self.assertEqual(expected, actual)


class RDSMovie(Base):  # type: ignore
    __tablename__ = 'movie'

    rk = Column(String(128), primary_key=True)
    name = Column(String(128))


class RDSActor(Base):  # type: ignore
    __tablename__ = 'actor'

    rk = Column(String(128), primary_key=True)
    name = Column(String(128))


class RDSMovieActor(Base):  # type: ignore
    __tablename__ = 'movie_actor'

    movie_rk = Column(String(128), ForeignKey('movie.rk'), primary_key=True)
    actor_rk = Column(String(128), ForeignKey('actor.rk'), primary_key=True)


class Actor(object):
    KEY_FORMAT = 'actor://{}'

    def __init__(self, name: str) -> None:
        self.name = name


class Movie(TableSerializable):
    KEY_FORMAT = 'movie://{}'

    def __init__(self,
                 name: str,
                 actors: Iterable[Actor]) -> None:
        self._name = name
        self._actors = actors
        self._record_iter = self._create_record_iterator()

    def create_next_record(self) -> Union[RDSModel, None]:
        try:
            return next(self._record_iter)
        except StopIteration:
            return None

    def _create_record_iterator(self) -> Iterator[RDSModel]:
        movie_record = RDSMovie(
            rk=Movie.KEY_FORMAT.format(self._name),
            name=self._name
        )
        yield movie_record

        for actor in self._actors:
            actor_record = RDSActor(
                rk=Actor.KEY_FORMAT.format(actor.name),
                name=actor.name
            )
            yield actor_record

            movie_actor_record = RDSMovieActor(
                movie_rk=Movie.KEY_FORMAT.format(self._name),
                actor_rk=Actor.KEY_FORMAT.format(actor.name)
            )
            yield movie_actor_record


if __name__ == '__main__':
    unittest.main()
