from __future__ import annotations

import json
from pathlib import Path
from typing import Any


class OpenLaneResultWriter:
    def __init__(self, jsonl_path: str | Path, aggregated_path: str | Path) -> None:
        self.jsonl_path = Path(jsonl_path)
        self.aggregated_path = Path(aggregated_path)

    def ensure_parent_dirs(self) -> None:
        self.jsonl_path.parent.mkdir(parents=True, exist_ok=True)
        self.aggregated_path.parent.mkdir(parents=True, exist_ok=True)

    def append_page(self, page: int, records: list[dict[str, Any]]) -> int:
        self.ensure_parent_dirs()
        written = 0
        with self.jsonl_path.open("a", encoding="utf-8") as fh:
            for record in records:
                payload = {
                    "page": page,
                    "record": record,
                }
                fh.write(json.dumps(payload, ensure_ascii=False) + "\n")
                written += 1
        return written

    def finalize(self, *, max_pages: int, completed_pages: list[int], failed_pages: list[int]) -> dict[str, Any]:
        self.ensure_parent_dirs()
        records: list[dict[str, Any]] = []
        if self.jsonl_path.exists():
            with self.jsonl_path.open("r", encoding="utf-8") as fh:
                for line in fh:
                    line = line.strip()
                    if not line:
                        continue
                    item = json.loads(line)
                    records.append(item)

        summary = {
            "max_pages": max_pages,
            "completed_pages": sorted(completed_pages),
            "failed_pages": sorted(failed_pages),
            "total_pages_completed": len(set(completed_pages)),
            "total_pages_failed": len(set(failed_pages)),
            "total_records": len(records),
            "items": records,
        }
        self.aggregated_path.write_text(
            json.dumps(summary, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        return summary