import os
import csv
import json
from datetime import datetime
from workflow.module import Module
from workflow.io_port import InPort, AnyType
# Codes Go Below:


class Writer(Module):
    """
    Generic Writer node.
    Saves any data to disk in the desired format.
    """

    def __init__(self):
        super().__init__("Writer")

        # Accept any data type
        self.addInPort(InPort("data", AnyType()))

    def run(self, context):
        data = context.get_asset_value("Writer.data")

        # Metadata from Model layer
        output_dir = self.get_meta("output_dir", "results")
        file_format = self.get_meta("format", "csv")

        os.makedirs(output_dir, exist_ok=True)
        ts = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"output_{ts}.{file_format}"
        path = os.path.join(output_dir, filename)

        # Save based on format
        if file_format == "csv":
            with open(path, "w", newline="", encoding="utf-8") as f:
                w = csv.writer(f)
                if isinstance(data, list):
                    if isinstance(data[0], dict):
                        w.writerow(data[0].keys())
                        for row in data:
                            w.writerow(row.values())
                    else:
                        for row in data:
                            w.writerow([row])
                else:
                    w.writerow([data])

        elif file_format == "json":
            with open(path, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)

        elif file_format == "txt":
            with open(path, "w", encoding="utf-8") as f:
                f.write(str(data))

        else:
            raise ValueError(f"Unsupported format: {file_format}")

        return {"path": path}