"""
Resumable chunked writer for Part 2 analysis results, one row per
image_id. Same chunk-then-consolidate pattern as ManifestWriter (see
manifest.py) so a crashed/interrupted analysis run can resume without
re-calling Gemini or re-computing embeddings for images already done.
"""
import glob
import os
from datetime import datetime, timezone

import pandas as pd

ANALYSIS_COLUMNS = [
    "image_id",
    "race_classification",
    "L", "a", "b", "h_star", "skin_tone_error",
    "clip_emb", "face_emb", "face_det_score", "face_bbox_xywh",
    "n_faces_detected", "embed_status",
    "error",
    "analyzed_at",
]


class AnalysisWriter:
    def __init__(self, output_dir, checkpoint_every=25):
        self.output_dir = output_dir
        self.chunk_dir = os.path.join(output_dir, "analysis_chunks")
        os.makedirs(self.chunk_dir, exist_ok=True)
        self.checkpoint_every = checkpoint_every
        self._buffer = []
        self._chunk_idx = len(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet")))
        self.done_image_ids = self._load_done_ids()

    def _load_done_ids(self):
        chunks = sorted(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet")))
        if not chunks:
            return set()
        dfs = [pd.read_parquet(c, columns=["image_id"]) for c in chunks]
        return set(pd.concat(dfs, ignore_index=True)["image_id"])

    def is_done(self, image_id):
        return image_id in self.done_image_ids

    def add_row(self, image_id, **fields):
        row = {"image_id": image_id, "analyzed_at": datetime.now(timezone.utc).isoformat()}
        row.update({col: fields.get(col) for col in ANALYSIS_COLUMNS if col not in row})
        self._buffer.append(row)
        self.done_image_ids.add(image_id)
        if len(self._buffer) >= self.checkpoint_every:
            self.flush()
        return image_id

    def flush(self):
        if not self._buffer:
            return
        df = pd.DataFrame(self._buffer, columns=ANALYSIS_COLUMNS)
        path = os.path.join(self.chunk_dir, f"chunk_{self._chunk_idx:05d}.parquet")
        df.to_parquet(path, index=False)
        self._chunk_idx += 1
        self._buffer = []

    def finalize(self, out_path=None):
        self.flush()
        out_path = out_path or os.path.join(self.output_dir, "analysis.parquet")
        chunks = sorted(glob.glob(os.path.join(self.chunk_dir, "chunk_*.parquet")))
        if not chunks:
            raise RuntimeError("No chunks to consolidate -- did analysis produce any rows?")
        df = pd.concat([pd.read_parquet(c) for c in chunks], ignore_index=True)
        df = df.drop_duplicates(subset=["image_id"], keep="last").reset_index(drop=True)
        df.to_parquet(out_path, index=False)
        return out_path, df