import csv
from pathlib import Path
from pydantic import BaseModel
from .utils import get_fields_for_model
from .types import UsageTypes


class CSVWriter:
    """CSV writer that logs usage and cost models in a structured format."""

    def __init__(self, name: str | Path):
        self.path = Path(name)
        self.path.parent.mkdir(parents=True, exist_ok=True)

        file_exists = self.path.exists()
        self._file = self.path.open("a", newline="")
        self._writer: csv.DictWriter | None = None
        self._header_written = file_exists

    def initialize(self, pricing: dict[str, float]) -> None:
        # CSV does not need pricing, but keep API parity with other writers
        pass

    def write(
        self,
        model: str,
        result: str,
        usage: UsageTypes,
        costs: BaseModel,
    ):
        row = {
            "Model": model,
            "Result": result,
            "Total Tokens": usage.total_tokens,
            "Input Tokens": usage.input_tokens,
            "Output Tokens": usage.output_tokens,
        }

        row.update(get_fields_for_model(costs))

        if self._writer is None:
            self._writer = csv.DictWriter(self._file, fieldnames=row.keys())

            if not self._header_written:
                self._writer.writeheader()
                self._header_written = True

        self._writer.writerow(row)
        self._file.flush()

    def close(self):
        if self._file:
            self._file.close()
            self._file = None
            self._writer = None
