from dataclasses import dataclass, fields, field
import csv
import os
from typing import List

@dataclass
class ImageGenerationRecord():
    """
    A data container for tracking and saving image generation metadata.

    This class serves as a structured record for a single image generation event.
    It includes functionality to serialize its own data to a CSV file, handling
    formatting (e.g., float precision) and file management (creation vs. appending).

    Attributes:
        image_gen_data_file_path (str): Path to the CSV file where records are saved.
        filename (str): Name of the generated image file.
        timestamp (str): Timestamp of generation.
        model (str): Name of the model used.
        device (str): Compute device used (e.g., 'cuda', 'cpu').
        dtype (str): Data type used (e.g., 'float16').
        prompt (str): The text prompt used for generation.
        seed (int): Random seed used.
        height (int): Image height.
        width (int): Image width.
        inf_steps (int): Number of inference steps.
        guidance_scale (float): Classifier-free guidance scale.
        total_generation_time (float): Total time taken for the batch/image.
        generation_time_per_image (float): Normalized time per single image.
        image_rating (float): Placeholder for user rating (defaults to -1).
        fields_to_not_save (list[str]): Internal list of field names to exclude from the CSV.
    """
    image_gen_data_file_path: str = ""
    filename: str = ""
    timestamp: str = ""
    model: str = ""
    device: str = ""
    dtype: str = ""
    prompt: str = ""
    seed: int = 0
    height: int = 0
    width: int = 0
    inf_steps: int = 0
    guidance_scale: float = 0
    total_generation_time: float = 0
    generation_time_per_image: float = 0
    image_rating: float = -1
    fields_to_not_save: list[str] = field(default_factory=list)


    def __post_init__(self):
        """
        Post-initialization hook to define excluded fields.

        Populates `fields_to_not_save` with technical fields that should not
        appear in the final CSV report (e.g., file paths or the exclusion list itself).
        """
        self.fields_to_not_save += ["image_gen_data_file_path", "fields_to_not_save", "path_to_csv"]


    def create_data_row(self):
        """
        Constructs a list of values for the CSV row based on current attributes.

        Iterates through the dataclass fields, filtering out those in `fields_to_not_save`.
        Floats are formatted to 2 decimal places.

        Returns:
            list: A list of values corresponding to the CSV columns.
        """
        row = []
        for field in fields(self):
            if field.name in self.fields_to_not_save:
                continue
            field_value = getattr(self, field.name)

            if isinstance(field_value, float):
                row.append(f"{field_value:.2f}")
            else:
                row.append(field_value)
        return row


    def append_data_to_csv(self):
        """
        Appends the current record to the CSV file.

        Opens the file at `image_gen_data_file_path` in append mode and writes
        the data row generated by `create_data_row()`.
        """
        with open(self.image_gen_data_file_path, "a", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(self.create_data_row())


    def create_header_row(self):
        """
        Generates the CSV header row based on field names.

        Returns:
            list[str]: A list of column names, excluding those in `fields_to_not_save`.
        """
        headers = []
        for header in fields(self):
            if header.name in self.fields_to_not_save:
                continue
            headers.append(header.name)
        return headers


    def create_new_stats_file(self):
        """
        Creates a new CSV file and writes the header row.

        Uses 'x' mode to open the file, which ensures a new file is created
        and raises an error if the file already exists (though `save_data()` guards against this).
        """
        with open(self.image_gen_data_file_path, "x", newline="") as f:
            writer = csv.writer(f)
            writer.writerow(self.create_header_row())


    def save_data(self):
        """
        Saves the record of the image generation stats to disk.

        Checks if the target CSV file exists:
        - If no: Creates the file and writes headers.
        - If yes: Appends the current data row.
        """
        if not os.path.exists(self.image_gen_data_file_path):
            self.create_new_stats_file()

        self.append_data_to_csv()
        
