from __future__ import annotations

import contextlib
import hashlib
import json
import os
import re
import secrets
import shutil
import time
import uuid
from dataclasses import dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, TYPE_CHECKING, Literal

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.parquet as pq

from .manifest import ArtifactEntry, RunManifest, RunSummary, TableEntry, manifest_from_json, manifest_to_json, utc_now_iso
from .policy import Policy

if TYPE_CHECKING:
    import pandas as pd

_SIMPLE_SEGMENT_RE = re.compile(r"^[A-Za-z0-9_-]+$")
_ARTIFACT_NAME_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
_LOCK_TIMEOUT_SECONDS = 10.0
_LOCK_POLL_SECONDS = 0.05
_ATOMIC_REPLACE_RETRIES = 20
_ATOMIC_REPLACE_POLL_SECONDS = 0.05
_LOCKED_STORAGE_MODES = {"compacted_session", "partitioned_by_run"}
_DEFAULT_COMPACTED_FILE_BUCKETS = 32


@dataclass(slots=True, frozen=True)
class StorageModePolicy:
    compacted_rows_per_run_max: int = 50
    compacted_min_run_count: int = 200

    def __post_init__(self) -> None:
        if self.compacted_rows_per_run_max <= 0:
            raise ValueError("compacted_rows_per_run_max must be positive")
        if self.compacted_min_run_count <= 0:
            raise ValueError("compacted_min_run_count must be positive")

    def to_dict(self) -> dict[str, int]:
        return {
            "compacted_rows_per_run_max": int(self.compacted_rows_per_run_max),
            "compacted_min_run_count": int(self.compacted_min_run_count),
        }


@dataclass(slots=True)
class RunHandle:
    pipeline_id: str
    run_id: str
    session_id: str
    session_run_index: int
    session_path: Path
    run_path: Path
    manifest: RunManifest


class DataHiveClient:
    """Filesystem-backed gatekeeper for canonical data hive run storage."""

    def __init__(
        self,
        workspace_root: str,
        policy: Policy | None = None,
        storage_mode_policy: StorageModePolicy | None = None,
    ) -> None:
        self._workspace_root = Path(workspace_root).expanduser().resolve()
        self._data_hive_root = self._workspace_root / "data_hive"
        self._data_hive_root.mkdir(parents=True, exist_ok=True)
        self._policy = policy or Policy()
        self._storage_mode_policy = storage_mode_policy or StorageModePolicy()

    def create_run(
        self,
        pipeline_id: str,
        pipeline_label: str | None = None,
        metadata: dict | None = None,
    ) -> RunHandle:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        self._require_write(pipeline_id, op="create_run", resource="run")

        now_utc = datetime.now(timezone.utc)
        session_id = _session_id_for_date(now_utc)
        session_path = self._session_dir(pipeline_id, now_utc, session_id)
        runs_dir = session_path / "runs"
        runs_dir.mkdir(parents=True, exist_ok=True)
        (session_path / "tables").mkdir(parents=True, exist_ok=True)

        run_id = _generate_run_id()
        run_path = runs_dir / run_id
        while run_path.exists():
            run_id = _generate_run_id()
            run_path = runs_dir / run_id

        run_run_index = self._next_session_run_index(session_path)
        (run_path / "artifacts").mkdir(parents=True, exist_ok=True)

        manifest = RunManifest(
            pipeline_id=pipeline_id,
            pipeline_label=pipeline_label,
            run_id=run_id,
            session_id=session_id,
            session_run_index=run_run_index,
            created_at=utc_now_iso(),
            status="staged",
            metadata=dict(metadata) if metadata is not None else None,
            tables={},
            artifacts={},
            semantic_refs=None,
        )

        run = RunHandle(
            pipeline_id=pipeline_id,
            run_id=run_id,
            session_id=session_id,
            session_run_index=run_run_index,
            session_path=session_path,
            run_path=run_path,
            manifest=manifest,
        )
        self._write_manifest(run)
        return run

    def open_run(self, pipeline_id: str, run_id: str) -> RunHandle:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        run_id = _sanitize_segment(run_id, field="run_id")
        self._require_read(pipeline_id, op="open_run", resource=run_id)
        self._require_write(pipeline_id, op="open_run", resource=run_id)

        run_path, session_path = self._locate_run_paths(pipeline_id, run_id)
        manifest_path = run_path / "manifest.json"
        if not manifest_path.exists():
            raise FileNotFoundError(f"manifest not found: {manifest_path}")

        manifest = _load_manifest_file(manifest_path)
        if manifest.pipeline_id != pipeline_id or manifest.run_id != run_id:
            raise ValueError("manifest pipeline/run mismatch")
        session_id = manifest.session_id or session_path.name
        session_run_index = int(manifest.session_run_index or 0)
        return RunHandle(
            pipeline_id=pipeline_id,
            run_id=run_id,
            session_id=session_id,
            session_run_index=session_run_index,
            session_path=session_path,
            run_path=run_path,
            manifest=manifest,
        )

    def write_table(
        self,
        run: RunHandle,
        table_name: str,
        data: pa.Table | pd.DataFrame,
        step_id: str | None = None,
        mode: Literal["append", "overwrite"] = "append",
        chunk_size: int | None = None,
        storage_mode: Literal["auto", "compacted_session", "partitioned_by_run"] = "auto",
        rows_per_run_estimate: int | None = None,
        run_count_estimate: int | None = None,
    ) -> TableEntry:
        table_name = _sanitize_segment(table_name, field="table_name")
        self._require_write(run.pipeline_id, op="write_table", resource=table_name)
        self._ensure_staged(run)

        table = _to_arrow_table(data)
        schema_fingerprint = _schema_fingerprint(table.schema)

        if storage_mode not in {"auto", "compacted_session", "partitioned_by_run"}:
            raise ValueError(f"invalid storage_mode: {storage_mode}")
        if rows_per_run_estimate is not None and rows_per_run_estimate <= 0:
            raise ValueError("rows_per_run_estimate must be positive when provided")
        if run_count_estimate is not None and run_count_estimate <= 0:
            raise ValueError("run_count_estimate must be positive when provided")

        table_root = run.session_path / "tables" / table_name
        with _run_lock(run.session_path / ".session.lock"):
            with _run_lock(run.run_path / ".lock"):
                previous = run.manifest.tables.get(table_name)
                if mode == "append":
                    if previous and previous.schema_fingerprint != schema_fingerprint:
                        raise ValueError(
                            "schema_fingerprint mismatch for append operation "
                            f"on table '{table_name}'"
                        )
                elif mode != "overwrite":
                    raise ValueError(f"invalid write mode: {mode}")

                resolved_mode, lock_reason = self._resolve_storage_mode(
                    run,
                    storage_mode=storage_mode,
                    table=table,
                    rows_per_run_estimate=rows_per_run_estimate,
                    run_count_estimate=run_count_estimate,
                )
                locked_mode, mode_lock_reason = self._lock_table_storage_mode(
                    run,
                    table_name=table_name,
                    requested_mode=resolved_mode,
                    reason=lock_reason,
                )
                self._assert_storage_mode_exclusive(table_root, locked_mode)

                if locked_mode == "partitioned_by_run":
                    table_dir = table_root / f"run_id={run.run_id}"
                    if mode == "overwrite" and table_dir.exists():
                        shutil.rmtree(table_dir)
                    table_dir.mkdir(parents=True, exist_ok=True)
                    _write_parquet_parts(table_dir, table, chunk_size=chunk_size)
                else:
                    table = _ensure_run_id_column(table, run.run_id)
                    table_dir = table_root / "_session_compacted"
                    if mode == "overwrite" and table_dir.exists():
                        shutil.rmtree(table_dir)
                    table_dir.mkdir(parents=True, exist_ok=True)
                    _write_compacted_table(
                        table_dir,
                        table,
                        run_id=run.run_id,
                        mode=mode,
                        max_files=_DEFAULT_COMPACTED_FILE_BUCKETS,
                    )

                n_rows, n_files = _parquet_dir_stats(table_dir)
                entry = TableEntry(
                    path=table_dir.relative_to(run.session_path).as_posix(),
                    format="parquet",
                    schema_fingerprint=schema_fingerprint,
                    n_rows=n_rows,
                    n_files=n_files,
                    file_hashes=None,
                    created_by_step=step_id,
                    created_at=utc_now_iso(),
                    locked_storage_mode=locked_mode,  # type: ignore[arg-type]
                    storage_mode_policy=self._storage_mode_policy.to_dict(),
                    mode_lock_reason=mode_lock_reason,
                )
                run.manifest.tables[table_name] = entry
                self._write_manifest(run)
                return entry

    def write_artifact(
        self,
        run: RunHandle,
        artifact_name: str,
        payload: bytes | str | os.PathLike[str],
        media_type: str | None = None,
        step_id: str | None = None,
    ) -> ArtifactEntry:
        artifact_name = _sanitize_artifact_name(artifact_name)
        self._require_write(run.pipeline_id, op="write_artifact", resource=artifact_name)
        self._ensure_staged(run)

        data = _artifact_payload_bytes(payload)
        digest = hashlib.sha256(data).hexdigest()

        artifact_path = run.run_path / "artifacts" / artifact_name
        with _run_lock(run.run_path / ".lock"):
            artifact_path.parent.mkdir(parents=True, exist_ok=True)
            _atomic_write_bytes(artifact_path, data)

            entry = ArtifactEntry(
                path=artifact_path.relative_to(run.run_path).as_posix(),
                media_type=media_type,
                sha256=digest,
                size_bytes=len(data),
                created_by_step=step_id,
                created_at=utc_now_iso(),
            )
            run.manifest.artifacts[artifact_name] = entry
            self._write_manifest(run)
            return entry

    def commit_run(self, run: RunHandle) -> None:
        self._require_write(run.pipeline_id, op="commit_run", resource=run.run_id)

        with _run_lock(run.run_path / ".lock"):
            self._ensure_staged(run)
            run.manifest.session_id = run.session_id
            run.manifest.session_run_index = int(run.session_run_index)
            self._refresh_table_entries(run)
            self._refresh_artifact_entries(run)
            run.manifest.status = "committed"
            self._write_manifest(run)

            # Serialize updates to shared "latest*" pointers at pipeline scope.
            with _run_lock(self._pipeline_dir(run.pipeline_id) / ".latest.lock"):
                latest_path = self._pipeline_dir(run.pipeline_id) / "latest.txt"
                _atomic_write_text(latest_path, run.run_id)
                latest_session_path = self._pipeline_dir(run.pipeline_id) / "latest_session.txt"
                _atomic_write_text(latest_session_path, run.session_id)

    def abort_run(self, run: RunHandle, reason: str | None = None) -> None:
        self._require_write(run.pipeline_id, op="abort_run", resource=run.run_id)

        with _run_lock(run.run_path / ".lock"):
            if run.manifest.status == "committed":
                raise RuntimeError("cannot abort a committed run")
            run.manifest.status = "aborted"
            if reason:
                merged_metadata = dict(run.manifest.metadata or {})
                merged_metadata["abort_reason"] = reason
                run.manifest.metadata = merged_metadata
            self._write_manifest(run)

    def list_runs(self, pipeline_id: str) -> list[RunSummary]:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        self._require_read(pipeline_id, op="list_runs", resource="runs")

        pipeline_dir = self._pipeline_dir(pipeline_id)
        if not pipeline_dir.exists():
            return []

        summaries: list[RunSummary] = []
        for manifest_path in pipeline_dir.rglob("manifest.json"):
            run_dir = manifest_path.parent
            if run_dir.parent.name != "runs":
                continue
            manifest = _load_manifest_file(manifest_path)
            summaries.append(
                RunSummary(
                    pipeline_id=manifest.pipeline_id,
                    pipeline_label=manifest.pipeline_label,
                    run_id=manifest.run_id,
                    created_at=manifest.created_at,
                    status=manifest.status,
                    session_id=manifest.session_id,
                    session_run_index=manifest.session_run_index,
                )
            )

        summaries.sort(
            key=lambda item: (
                item.created_at,
                item.run_id,
            ),
            reverse=True,
        )
        return summaries

    def load_manifest(self, pipeline_id: str, run_id: str) -> RunManifest:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        run_id = _sanitize_segment(run_id, field="run_id")
        self._require_read(pipeline_id, op="load_manifest", resource=run_id)

        run_path, _ = self._locate_run_paths(pipeline_id, run_id)
        manifest_path = run_path / "manifest.json"
        if not manifest_path.exists():
            raise FileNotFoundError(f"manifest not found: {manifest_path}")
        return _load_manifest_file(manifest_path)

    def read_table(
        self,
        pipeline_id: str,
        run_id: str,
        table_name: str,
        *,
        read_scope: Literal["run", "session"] = "run",
    ) -> pa.Table:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        run_id = _sanitize_segment(run_id, field="run_id")
        table_name = _sanitize_segment(table_name, field="table_name")
        self._require_read(pipeline_id, op="read_table", resource=table_name)
        if read_scope not in {"run", "session"}:
            raise ValueError(f"invalid read_scope: {read_scope}")

        run = self.open_run(pipeline_id, run_id)
        manifest = run.manifest
        if table_name not in manifest.tables:
            raise KeyError(f"table not found in manifest: {table_name}")

        table_entry = manifest.tables[table_name]
        table_path = _resolve_run_relative_path(
            run.session_path,
            table_entry.path,
            field=f"table path for '{table_name}'",
        )
        if not table_path.exists():
            raise FileNotFoundError(f"table path not found: {table_path}")

        locked_mode = table_entry.locked_storage_mode
        if locked_mode is None:
            locked_mode = _infer_storage_mode_from_path(table_entry.path)

        if read_scope == "run":
            if locked_mode == "compacted_session":
                session_table = _read_parquet_path(table_path)
                if "run_id" not in session_table.column_names:
                    raise RuntimeError(
                        "compacted_session table is missing required 'run_id' column: "
                        f"{table_path}"
                    )
                run_mask = pc.equal(session_table["run_id"], pa.scalar(run_id))
                return session_table.filter(run_mask)
            return _read_parquet_path(table_path)

        table_root = run.session_path / "tables" / table_name
        if locked_mode == "compacted_session":
            compacted_path = table_root / "_session_compacted"
            if not compacted_path.exists():
                raise FileNotFoundError(f"compacted session table path not found: {compacted_path}")
            return _read_parquet_path(compacted_path)

        parquet_files = sorted(table_root.glob("run_id=*/part-*.parquet"))
        if not parquet_files:
            raise FileNotFoundError(f"no run-partition parquet files found in table path: {table_root}")
        return _read_parquet_files(parquet_files)

    def open_artifact(self, pipeline_id: str, run_id: str, artifact_name: str) -> bytes:
        pipeline_id = _sanitize_segment(pipeline_id, field="pipeline_id")
        run_id = _sanitize_segment(run_id, field="run_id")
        artifact_name = _sanitize_artifact_name(artifact_name)
        self._require_read(pipeline_id, op="open_artifact", resource=artifact_name)

        run = self.open_run(pipeline_id, run_id)
        manifest = run.manifest
        if artifact_name not in manifest.artifacts:
            raise KeyError(f"artifact not found in manifest: {artifact_name}")

        artifact_entry = manifest.artifacts[artifact_name]
        artifact_path = _resolve_run_relative_path(
            run.run_path,
            artifact_entry.path,
            field=f"artifact path for '{artifact_name}'",
        )
        if not artifact_path.exists():
            raise FileNotFoundError(f"artifact path not found: {artifact_path}")
        return artifact_path.read_bytes()

    def _pipeline_dir(self, pipeline_id: str) -> Path:
        return self._data_hive_root / pipeline_id

    def _session_dir(self, pipeline_id: str, now_utc: datetime, session_id: str) -> Path:
        return (
            self._pipeline_dir(pipeline_id)
            / now_utc.strftime("%Y")
            / now_utc.strftime("%m")
            / now_utc.strftime("%d")
            / session_id
        )

    @staticmethod
    def _next_session_run_index(session_path: Path) -> int:
        runs_dir = session_path / "runs"
        if not runs_dir.exists():
            return 0
        count = 0
        for run_dir in runs_dir.iterdir():
            if run_dir.is_dir() and (run_dir / "manifest.json").exists():
                count += 1
        return count

    def _locate_run_paths(self, pipeline_id: str, run_id: str) -> tuple[Path, Path]:
        pipeline_dir = self._pipeline_dir(pipeline_id)

        legacy_run_path = pipeline_dir / "runs" / run_id
        legacy_manifest = legacy_run_path / "manifest.json"
        if legacy_manifest.exists():
            return legacy_run_path, pipeline_dir

        for manifest_path in pipeline_dir.rglob("manifest.json"):
            run_path = manifest_path.parent
            if run_path.name != run_id:
                continue
            if run_path.parent.name != "runs":
                continue
            session_path = run_path.parent.parent
            return run_path, session_path
        raise FileNotFoundError(f"run not found: pipeline={pipeline_id} run={run_id}")

    def _resolve_storage_mode(
        self,
        run: RunHandle,
        *,
        storage_mode: Literal["auto", "compacted_session", "partitioned_by_run"],
        table: pa.Table,
        rows_per_run_estimate: int | None,
        run_count_estimate: int | None,
    ) -> tuple[str, str]:
        if storage_mode in _LOCKED_STORAGE_MODES:
            return storage_mode, f"explicit:{storage_mode}"

        rows_estimate = int(rows_per_run_estimate if rows_per_run_estimate is not None else table.num_rows)
        run_count = (
            int(run_count_estimate)
            if run_count_estimate is not None
            else self._next_session_run_index(run.session_path)
        )
        if run_count <= 0:
            run_count = 1

        policy = self._storage_mode_policy
        use_compacted = (
            rows_estimate <= int(policy.compacted_rows_per_run_max)
            and run_count >= int(policy.compacted_min_run_count)
        )
        resolved = "compacted_session" if use_compacted else "partitioned_by_run"
        reason = (
            "auto:"
            f"rows_per_run_estimate={rows_estimate},"
            f"run_count={run_count},"
            f"threshold_rows={policy.compacted_rows_per_run_max},"
            f"threshold_runs={policy.compacted_min_run_count}"
        )
        return resolved, reason

    def _lock_table_storage_mode(
        self,
        run: RunHandle,
        *,
        table_name: str,
        requested_mode: str,
        reason: str,
    ) -> tuple[str, str]:
        if requested_mode not in _LOCKED_STORAGE_MODES:
            raise ValueError(f"invalid requested_mode: {requested_mode}")

        session_manifest = self._load_session_manifest(run)
        tables = dict(session_manifest.get("tables") or {})
        table_lock = dict(tables.get(table_name) or {})
        locked_mode = table_lock.get("locked_storage_mode")
        lock_reason = table_lock.get("mode_lock_reason")

        if locked_mode is None:
            inferred_mode = _detect_storage_mode_from_filesystem(run.session_path / "tables" / table_name)
            if inferred_mode is not None:
                locked_mode = inferred_mode
                lock_reason = "inferred_from_filesystem"

        if locked_mode is not None:
            locked_mode = str(locked_mode)
            if locked_mode not in _LOCKED_STORAGE_MODES:
                raise ValueError(f"invalid locked storage mode found in session manifest: {locked_mode}")
            if locked_mode != requested_mode:
                raise ValueError(
                    "storage mode lock violation for "
                    f"table '{table_name}' in session '{run.session_id}': "
                    f"locked='{locked_mode}', attempted='{requested_mode}'"
                )
        else:
            locked_mode = requested_mode
            lock_reason = reason

        tables[table_name] = {
            "locked_storage_mode": locked_mode,
            "mode_lock_reason": str(lock_reason) if lock_reason is not None else reason,
            "storage_mode_policy": self._storage_mode_policy.to_dict(),
            "locked_at": table_lock.get("locked_at", utc_now_iso()),
        }
        session_manifest["tables"] = tables
        session_manifest["updated_at"] = utc_now_iso()
        self._write_session_manifest(run.session_path, session_manifest)
        return locked_mode, str(tables[table_name]["mode_lock_reason"])

    @staticmethod
    def _assert_storage_mode_exclusive(table_root: Path, locked_mode: str) -> None:
        compacted_dir = table_root / "_session_compacted"
        run_partition_dirs = [
            candidate
            for candidate in table_root.glob("run_id=*")
            if candidate.is_dir()
        ] if table_root.exists() else []

        if locked_mode == "compacted_session":
            if run_partition_dirs:
                raise RuntimeError(
                    "mixed storage mode detected: compacted_session is locked but "
                    f"run_id partitions exist under {table_root}"
                )
            return
        if locked_mode == "partitioned_by_run":
            if compacted_dir.exists():
                raise RuntimeError(
                    "mixed storage mode detected: partitioned_by_run is locked but "
                    f"_session_compacted exists under {table_root}"
                )
            return
        raise ValueError(f"invalid locked_mode: {locked_mode}")

    def _load_session_manifest(self, run: RunHandle) -> dict[str, Any]:
        manifest_path = run.session_path / "session_manifest.json"
        if not manifest_path.exists():
            now = utc_now_iso()
            return {
                "schema_version": "1.0",
                "pipeline_id": run.pipeline_id,
                "session_id": run.session_id,
                "created_at": now,
                "updated_at": now,
                "tables": {},
            }

        payload = json.loads(manifest_path.read_text(encoding="utf-8"))
        if str(payload.get("pipeline_id")) != run.pipeline_id:
            raise ValueError(
                f"session_manifest pipeline_id mismatch: {payload.get('pipeline_id')} != {run.pipeline_id}"
            )
        if str(payload.get("session_id")) != run.session_id:
            raise ValueError(
                f"session_manifest session_id mismatch: {payload.get('session_id')} != {run.session_id}"
            )
        payload.setdefault("schema_version", "1.0")
        payload.setdefault("tables", {})
        payload.setdefault("created_at", utc_now_iso())
        payload.setdefault("updated_at", utc_now_iso())
        return payload

    @staticmethod
    def _write_session_manifest(session_path: Path, payload: dict[str, Any]) -> None:
        manifest_path = session_path / "session_manifest.json"
        _atomic_write_text(manifest_path, json.dumps(payload, indent=2, sort_keys=True))

    def _require_read(self, pipeline_id: str, op: str, resource: str) -> None:
        if not self._policy.allow_read(pipeline_id, op, resource):
            raise PermissionError(f"policy denied read op={op} resource={resource}")

    def _require_write(self, pipeline_id: str, op: str, resource: str) -> None:
        if not self._policy.allow_write(pipeline_id, op, resource):
            raise PermissionError(f"policy denied write op={op} resource={resource}")

    @staticmethod
    def _ensure_staged(run: RunHandle) -> None:
        if run.manifest.status != "staged":
            raise RuntimeError(f"run is not staged (status={run.manifest.status})")

    @staticmethod
    def _write_manifest(run: RunHandle) -> None:
        manifest_path = run.run_path / "manifest.json"
        _atomic_write_text(manifest_path, manifest_to_json(run.manifest))

    @staticmethod
    def _refresh_table_entries(run: RunHandle) -> None:
        for table_name, entry in run.manifest.tables.items():
            table_dir = _resolve_run_relative_path(
                run.session_path,
                entry.path,
                field=f"table path for '{table_name}'",
            )
            parquet_files = sorted(table_dir.glob("part-*.parquet"))
            total_rows = 0
            hashes: list[str] = []
            for parquet_file in parquet_files:
                metadata = pq.read_metadata(parquet_file)
                total_rows += int(metadata.num_rows)
                hashes.append(_sha256_file(parquet_file))
            entry.n_rows = total_rows
            entry.n_files = len(parquet_files)
            entry.file_hashes = hashes

    @staticmethod
    def _refresh_artifact_entries(run: RunHandle) -> None:
        for artifact_name, entry in run.manifest.artifacts.items():
            artifact_path = _resolve_run_relative_path(
                run.run_path,
                entry.path,
                field=f"artifact path for '{artifact_name}'",
            )
            if not artifact_path.exists():
                raise FileNotFoundError(f"artifact missing during commit: {artifact_path}")
            entry.sha256 = _sha256_file(artifact_path)
            entry.size_bytes = artifact_path.stat().st_size


def _generate_run_id() -> str:
    if hasattr(uuid, "uuid7"):
        return str(uuid.uuid7())
    timestamp = time.strftime("%Y%m%d-%H%M%SZ", time.gmtime())
    suffix = secrets.token_hex(4)
    return f"{timestamp}_{suffix}"


def _session_id_for_date(now_utc: datetime) -> str:
    return f"session-{now_utc.strftime('%Y%m%d')}"


def _sanitize_segment(name: str, *, field: str) -> str:
    if not isinstance(name, str) or not name:
        raise ValueError(f"{field} must be a non-empty string")
    if name.startswith("."):
        raise ValueError(f"{field} cannot start with '.'")
    if "/" in name or "\\" in name or ".." in name:
        raise ValueError(f"{field} cannot contain path separators or '..'")
    if not _SIMPLE_SEGMENT_RE.fullmatch(name):
        raise ValueError(f"{field} must match [A-Za-z0-9_-]+")
    return name


def _sanitize_artifact_name(name: str) -> str:
    if not isinstance(name, str) or not name:
        raise ValueError("artifact_name must be a non-empty string")
    if name.startswith("."):
        raise ValueError("artifact_name cannot start with '.'")
    if "/" in name or "\\" in name or ".." in name:
        raise ValueError("artifact_name cannot contain path separators or '..'")
    if not _ARTIFACT_NAME_RE.fullmatch(name):
        raise ValueError("artifact_name contains unsupported characters")
    return name


def _schema_fingerprint(schema: pa.Schema) -> str:
    payload = schema.serialize().to_pybytes()
    return hashlib.sha256(payload).hexdigest()


def _write_parquet_parts(table_dir: Path, table: pa.Table, chunk_size: int | None) -> list[Path]:
    if chunk_size is not None and chunk_size <= 0:
        raise ValueError("chunk_size must be positive when provided")

    start_index = _next_part_index(table_dir)
    written_files: list[Path] = []

    if table.num_rows == 0:
        part_path = table_dir / f"part-{start_index:04d}.parquet"
        pq.write_table(table, part_path)
        written_files.append(part_path)
        return written_files

    rows_per_file = chunk_size or table.num_rows
    index = start_index
    for offset in range(0, table.num_rows, rows_per_file):
        chunk = table.slice(offset, rows_per_file)
        part_path = table_dir / f"part-{index:04d}.parquet"
        pq.write_table(chunk, part_path)
        written_files.append(part_path)
        index += 1

    return written_files


def _write_compacted_table(
    table_dir: Path,
    table: pa.Table,
    *,
    run_id: str,
    mode: Literal["append", "overwrite"],
    max_files: int,
) -> Path:
    if max_files <= 0:
        raise ValueError("max_files must be positive")

    if mode == "overwrite":
        bucket_index = 0
    else:
        bucket_index = int(hashlib.sha256(run_id.encode("utf-8")).hexdigest(), 16) % max_files

    part_path = table_dir / f"part-{bucket_index:04d}.parquet"
    if mode == "append" and part_path.exists():
        existing = pq.ParquetFile(part_path).read()
        table = pa.concat_tables([existing, table])
    pq.write_table(table, part_path)
    return part_path


def _next_part_index(table_dir: Path) -> int:
    current_max = -1
    for parquet_file in table_dir.glob("part-*.parquet"):
        suffix = parquet_file.stem.removeprefix("part-")
        if suffix.isdigit():
            current_max = max(current_max, int(suffix))
    return current_max + 1


def _parquet_dir_stats(table_dir: Path) -> tuple[int, int]:
    parquet_files = sorted(table_dir.glob("part-*.parquet"))
    total_rows = 0
    for parquet_file in parquet_files:
        metadata = pq.read_metadata(parquet_file)
        total_rows += int(metadata.num_rows)
    return total_rows, len(parquet_files)


def _read_parquet_files(parquet_files: list[Path]) -> pa.Table:
    if not parquet_files:
        raise FileNotFoundError("no parquet files provided")
    tables = [pq.ParquetFile(parquet_file).read() for parquet_file in parquet_files]
    if len(tables) == 1:
        return tables[0]
    return pa.concat_tables(tables)


def _read_parquet_path(table_path: Path) -> pa.Table:
    if not table_path.exists():
        raise FileNotFoundError(f"table path not found: {table_path}")
    if table_path.is_dir():
        parquet_files = sorted(table_path.glob("part-*.parquet"))
        if not parquet_files:
            raise FileNotFoundError(f"no parquet parts found in table path: {table_path}")
        return _read_parquet_files(parquet_files)
    return pq.ParquetFile(table_path).read()


def _ensure_run_id_column(table: pa.Table, run_id: str) -> pa.Table:
    if "run_id" in table.column_names:
        return table
    run_id_values = pa.array([run_id] * int(table.num_rows), type=pa.string())
    return table.append_column("run_id", run_id_values)


def _infer_storage_mode_from_path(table_path: str) -> str:
    normalized = table_path.replace("\\", "/")
    if "/_session_compacted/" in f"/{normalized}/" or normalized.endswith("/_session_compacted"):
        return "compacted_session"
    if "/run_id=" in f"/{normalized}/":
        return "partitioned_by_run"
    raise ValueError(f"unable to infer storage mode from table path: {table_path}")


def _detect_storage_mode_from_filesystem(table_root: Path) -> str | None:
    if not table_root.exists():
        return None
    has_compacted = (table_root / "_session_compacted").exists()
    has_run_partitions = any(candidate.is_dir() for candidate in table_root.glob("run_id=*"))
    if has_compacted and has_run_partitions:
        raise RuntimeError(f"mixed storage layout detected under {table_root}")
    if has_compacted:
        return "compacted_session"
    if has_run_partitions:
        return "partitioned_by_run"
    return None


def _to_arrow_table(data: pa.Table | pd.DataFrame) -> pa.Table:
    if isinstance(data, pa.Table):
        return data

    try:
        import pandas as pd  # type: ignore[import-not-found]
    except ImportError as exc:
        raise TypeError("pandas is not available; pass a pyarrow.Table instead") from exc

    if isinstance(data, pd.DataFrame):
        return pa.Table.from_pandas(data, preserve_index=False)

    raise TypeError("data must be a pyarrow.Table or pandas.DataFrame")


def _artifact_payload_bytes(payload: bytes | str | os.PathLike[str]) -> bytes:
    if isinstance(payload, bytes):
        return payload
    path = Path(payload)
    if not path.exists():
        raise FileNotFoundError(f"artifact payload path does not exist: {path}")
    return path.read_bytes()


def _sha256_file(path: Path) -> str:
    hasher = hashlib.sha256()
    with path.open("rb") as file_obj:
        while True:
            block = file_obj.read(1024 * 1024)
            if not block:
                break
            hasher.update(block)
    return hasher.hexdigest()


@contextlib.contextmanager
def _run_lock(lock_path: Path):
    deadline = time.monotonic() + _LOCK_TIMEOUT_SECONDS
    fd: int | None = None

    while fd is None:
        try:
            fd = os.open(str(lock_path), os.O_CREAT | os.O_EXCL | os.O_WRONLY)
        except FileExistsError:
            if time.monotonic() >= deadline:
                raise TimeoutError(f"timed out acquiring run lock: {lock_path}")
            time.sleep(_LOCK_POLL_SECONDS)

    try:
        os.write(fd, f"pid={os.getpid()}\n".encode("utf-8"))
        yield
    finally:
        os.close(fd)
        with contextlib.suppress(FileNotFoundError):
            lock_path.unlink()


def _atomic_write_text(path: Path, text: str) -> None:
    _atomic_write_bytes(path, text.encode("utf-8"))


def _atomic_write_bytes(path: Path, payload: bytes) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tmp_path = path.with_name(f".{path.name}.tmp-{secrets.token_hex(6)}")
    try:
        with tmp_path.open("wb") as file_obj:
            file_obj.write(payload)
            file_obj.flush()
            os.fsync(file_obj.fileno())
        for attempt in range(_ATOMIC_REPLACE_RETRIES):
            try:
                os.replace(tmp_path, path)
                break
            except PermissionError:
                if attempt >= (_ATOMIC_REPLACE_RETRIES - 1):
                    raise
                time.sleep(_ATOMIC_REPLACE_POLL_SECONDS * (attempt + 1))
    finally:
        with contextlib.suppress(FileNotFoundError):
            tmp_path.unlink()


def _load_manifest_file(manifest_path: Path) -> RunManifest:
    return manifest_from_json(manifest_path.read_text(encoding="utf-8"))


def _resolve_run_relative_path(run_path: Path, relative_path: str, *, field: str) -> Path:
    relative = Path(relative_path)
    if relative.is_absolute():
        raise ValueError(f"{field} must be relative")
    if any(part in {"", ".", ".."} for part in relative.parts):
        raise ValueError(f"{field} contains invalid path segments")

    resolved_run = run_path.resolve()
    resolved_target = (run_path / relative).resolve()
    if not resolved_target.is_relative_to(resolved_run):
        raise ValueError(f"{field} escapes run directory")
    return resolved_target
