"""
CSV logger for training metrics.

Provides :class:`DynamicCSVLogger`, a lightweight logger that writes metrics
to a CSV file and automatically handles new columns added mid-training.
"""
import csv
from pathlib import Path
from typing import List, Dict, Optional, Union
import os


class DynamicCSVLogger:
    """CSV logger that handles dynamically discovered metric columns.

    Columns do not need to be declared upfront.  When a new key appears in a
    :meth:`log` call the file is rewritten with the updated header so all
    previous rows remain valid (missing values are left blank).

    The parent directory of *log_path* is created automatically.

    Args:
        log_path: Destination CSV file.  If the file already exists the logger
            appends to it; otherwise it is created fresh.

    Example:
        .. code-block:: python

            from pathlib import Path
            from investigation.logger.logger import DynamicCSVLogger

            logger = DynamicCSVLogger(Path(".investigation/demo/history.csv"))
            for step in range(100):
                logger.log({"global_step": step, "loss": 1 / (step + 1)})
            logger.close()
    """
    log_path: Path

    def __init__(self, log_path: Union[str, Path]):
        self.log_path = Path(log_path)
        self.log_path.parent.mkdir(parents=True, exist_ok=True)

        self.fieldnames: List[str] = []
        self.rows: List[Dict] = []
        self._file = None
        self._writer = None

        # 1. 启动时加载所有历史数据到内存
        self._load_all_history()

        # 2. 根据是否有数据，决定打开模式 ('a' 或 'w')
        mode = 'a' if self.fieldnames else 'w'
        self._file = open(self.log_path, mode, newline='')

        if self.fieldnames:
            self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)

    def _load_all_history(self):
        """Load all existing rows from disk into :attr:`rows` and :attr:`fieldnames`."""
        if not self.log_path.exists() or self.log_path.stat().st_size == 0:
            return

        with open(self.log_path, 'r', newline='') as f:
            reader = csv.DictReader(f)
            if reader.fieldnames:
                self.fieldnames = list(reader.fieldnames)
            # 读入所有行，注意：CSV读入的值默认为字符串
            self.rows = list(reader)

    def resync(self, current_step: int, step_key: str):
        """Truncate the log to align with a loaded checkpoint.

        When training resumes from a checkpoint at *current_step*, any rows
        logged beyond that step represent a future that no longer exists.
        ``resync`` removes those rows and rewrites the file so the log is
        consistent with the checkpoint state.

        Args:
            current_step: The step value of the loaded checkpoint.  All rows
                with ``step_key > current_step`` will be discarded.
            step_key: Column name used as the step axis (e.g. ``"global_step"``).
        """
        if not self.rows:
            return

        # 检查最后一条记录的 step
        try:
            last_recorded_step = int(float(self.rows[-1].get(step_key, -1)))
        except (ValueError, TypeError):
            # 如果解析失败，说明数据可能有问题，不做截断，安全返回
            return

        # 如果 Log 的记录比 Checkpoint 还旧（或相等），说明没有“未来数据”，无需操作
        if last_recorded_step <= current_step:
            return

        print(
            f"Logger: Checkpoint step is {current_step}, but log has future data (up to {last_recorded_step}). Truncating log..."
        )

        # --- 执行截断逻辑 ---
        valid_rows = []
        for row in self.rows:
            try:
                # 兼容 int 和 float 格式的字符串转换
                row_step = int(float(row.get(step_key, -1)))
                if row_step <= current_step:
                    valid_rows.append(row)
            except (ValueError, TypeError):
                # 如果某一行没有 step 或者格式错误，保留或丢弃视具体需求。
                # 这里选择保留以防止意外丢失非 step 数据
                valid_rows.append(row)

        self.rows = valid_rows

        # --- 强制重写文件 ---
        # 即使 fieldnames 没变，内容变少了，也需要重写文件
        self._rewrite_file()
        print(f"Logger: Truncated log. Kept {len(self.rows)} rows.")

    def log(self, metrics: Dict):
        """Write one row of metrics to the CSV.

        If *metrics* contains keys not seen before, the file is rewritten with
        an updated header that includes the new columns.

        Args:
            metrics: Mapping of metric name to scalar value for this step.
        """
        # Detect new fields
        new_fields = [k for k in metrics if k not in self.fieldnames]

        if new_fields or self._writer is None:
            self._update_fieldnames(new_fields)

        self.rows.append(metrics)
        self._writer.writerow(metrics)
        self._file.flush()

    def _update_fieldnames(self, new_fields: List[str]):
        """Rewrite the CSV with an extended header to accommodate *new_fields*.

        Args:
            new_fields: Column names that are not yet in :attr:`fieldnames`.
        """
        self.fieldnames.extend(new_fields)
        self._rewrite_file()

    def _rewrite_file(self):
        """Atomically rewrite the CSV file with current rows and fieldnames."""
        if self._file:
            self._file.close()

        with open(self.log_path, 'w', newline='') as f:
            writer = csv.DictWriter(f, fieldnames=self.fieldnames)
            writer.writeheader()
            writer.writerows(self.rows)

        self._file = open(self.log_path, 'a', newline='')
        self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)

    def close(self):
        """Flush and close the underlying file handle."""
        self._file.close()
        if self._file:
            self._file.close()


import pandas as pd
from pathlib import Path
from typing import Union, Dict, List

import fcntl


class GPFSPandasLogger:
    """Context-manager logger for GPFS parallel file systems.

    Designed for multi-process training on cluster file systems (e.g. GPFS /
    Lustre) where multiple workers may write to the same CSV concurrently.
    Uses ``fcntl`` advisory file locking to serialise writes.

    Use as a context manager — the lock is acquired on ``__enter__`` and
    released (with a full file rewrite) on ``__exit__``:

    .. code-block:: python

        from investigation.logger.logger import GPFSPandasLogger

        logger = GPFSPandasLogger(".investigation/demo/history.csv")
        with logger:
            logger.update(global_step=100, metrics={"loss": 0.42})

    .. note::
        ``fcntl`` is POSIX-only and not available on Windows.
    """

    _df: pd.DataFrame

    def __init__(self, log_path: Union[str, Path]):
        self.log_path = Path(log_path)
        self.log_path.parent.mkdir(parents=True, exist_ok=True)

        # 锁文件路径 (例如 log.csv -> log.csv.lock)
        self.lock_path = self.log_path.with_suffix(self.log_path.suffix + ".lock")

        # 内部状态
        self._lock_file = None  # 持有文件句柄

    def __enter__(self):
        """Acquire an exclusive file lock and load existing CSV data into memory."""
        # 1. 准备锁文件
        if not self.lock_path.parent.exists():
            self.lock_path.parent.mkdir(parents=True, exist_ok=True)

        # 使用 'a' (append) 模式打开，避免 'w' 模式可能在其他进程持有锁时截断文件
        self._lock_file = open(self.lock_path, 'a')

        # 2. 获取锁 (对应你提供的 global_io_lock 逻辑)
        # print(f"[IO Lock] Trying to acquire lock at {self.lock_path}...", flush=True)

        # LOCK_EX: 排他锁, 默认阻塞
        fcntl.lockf(self._lock_file, fcntl.LOCK_EX)

        # 3. 拿到锁后，安全的读取数据
        # print(f"[IO Lock] Acquired. Reading data...", flush=True)
        if self.log_path.exists() and self.log_path.stat().st_size > 0:
            try:
                self._df = pd.read_csv(self.log_path)
            except pd.errors.EmptyDataError:
                self._df = pd.DataFrame()
        else:
            self._df = pd.DataFrame()

        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Write the in-memory DataFrame back to disk, then release the lock."""
        try:
            # 只有在代码块没有抛出异常时才保存
            if exc_type is None and self._df is not None:
                # 排序优化可读性 (可选)
                if 'global_step' in self._df.columns:
                    self._df['global_step'] = pd.to_numeric(self._df['global_step'], errors='coerce')
                    self._df = self._df.sort_values('global_step')

                # --- 修改开始 ---
                # 使用 'w' 模式打开文件，并将文件句柄 f 传给 to_csv
                with open(self.log_path, 'w') as f:
                    self._df.to_csv(f, index=False)

                    # 1. 强制将 Python 内部缓冲区的数据写入 OS 内核缓冲区
                    f.flush()

                    # 2. 强制将 OS 内核缓冲区的数据落盘 (GPFS 需要这一步)
                    os.fsync(f.fileno())
                # --- 修改结束 ---
                # with 语句结束时会自动 safe close，不需要额外操作

        finally:
            # 释放锁逻辑
            # print(f"[IO Lock] Releasing lock...", flush=True)
            if self._lock_file:
                fcntl.lockf(self._lock_file, fcntl.LOCK_UN)
                self._lock_file.close()
                self._lock_file = None
            self._df = None

    def update(self, global_step: int, metrics: Dict):
        """Update or append a row in the in-memory DataFrame (no I/O).

        If a row with *global_step* already exists it is updated in-place;
        otherwise a new row is appended.  The file is only written when the
        context manager exits.

        Args:
            global_step: Step index for this metrics snapshot.
            metrics: Mapping of metric name to scalar value.

        Raises:
            RuntimeError: If called outside a ``with`` block.
        """
        if self._df is None:
            raise RuntimeError("Logger.update() must be called within a 'with' statement.")

        if self._df.empty:
            self._df = pd.DataFrame([{
                **metrics,
                'global_step': global_step,
            }])
            return

        # 查找匹配行
        mask = (self._df['global_step'] == global_step)

        if mask.any():
            # Update existing
            idx = self._df.index[mask][0]
            for k, v in metrics.items():
                self._df.at[idx, k] = v
        else:
            # Append new
            new_row = pd.DataFrame([{
                **metrics,
                'global_step': global_step,
            }])
            self._df = pd.concat([self._df, new_row], ignore_index=True)

    def get_df(self):
        return self._df
