import io
from typing import Type, List, Callable

import matplotlib.pyplot
import pandas as pd
from PIL import Image
from networkx import DiGraph
from prefect import Flow
from pydantic import BaseModel
from pydot import Dot

from feature_forge.domain.entities.dataset import Dataset
from feature_forge.domain.entities.entity_model import EntityModel


class Recipe(BaseModel):
    dataset: Dataset
    entity: Type[EntityModel]
    flow: Callable[[pd.DataFrame], pd.DataFrame]
    graph: Dot
    requested_features: List[str]
    all_features: List[str]

    class Config:
        arbitrary_types_allowed = True
        frozen = True

    def __str__(self):
        return self.__repr__()

    def __repr_args__(self):
        return [(a, v) if a not in ['dataset']
                else (a, self.dataset.__class__.__name__)
                for a, v in super().__repr_args__()]

    def show(self):
        image = Image.open(io.BytesIO(self.graph.create_png()))
        image.show()
