"""Durable, date-partitioned Parquet archive and acknowledged upload queue.

Each mutation replaces a daily partition atomically. A process/thread lock serializes
readers and writers; it is never held during an HTTP request.
"""

from contextlib import contextmanager
from datetime import datetime, timedelta, timezone
import fcntl
import json
import os
from pathlib import Path
import tempfile
import threading
import time
from zoneinfo import ZoneInfo

import pyarrow as pa
import pyarrow.parquet as pq

from .config import retention_cutoff, utc_text
from .coverage import daily_coverage, milliseconds, sample_coverage

NUMBERS = (
    "production_w", "production_ac_w", "consumption_w", "grid_w",
    "grid_import_w", "grid_export_w", "production_ac_total_wh",
    "grid_import_total_wh", "grid_export_total_wh",
)
SCHEMA = pa.schema([
    ("id", pa.string()), ("source_id", pa.string()), ("observed_at", pa.string()),
    ("observed_ms", pa.int64()), ("recorded_ns", pa.int64()),
    ("uploaded_at", pa.string()), ("payload", pa.string()),
    *((field, pa.float64()) for field in NUMBERS),
    ("boiler_enabled", pa.bool_()), ("boiler_reason", pa.string()),
    ("record_kind", pa.string()),
    ("interval_energy_wh", pa.struct([(field, pa.float64()) for field in ("production", "consumption", "self_consumption", "bought", "sold")])),
    ("power_ms", pa.int64()), ("power_ok", pa.bool_()),
    ("meter_ms", pa.int64()), ("meter_ok", pa.bool_()),
    ("boiler_ms", pa.int64()), ("boiler_ok", pa.bool_()),
], metadata={b"solar_relay_schema": b"1"})
COVERAGE_COLUMNS = ["observed_ms", "power_ms", "power_ok", "meter_ms", "meter_ok", "boiler_ms", "boiler_ok"]


class StoreError(OSError):
    """An archive could not be read or updated safely."""


class Store:
    def __init__(self, path: Path, retention_days=30, *, timezone_name="Europe/Zurich"):
        self.path = Path(path)
        if self.path.is_file():
            raise StoreError("collector.archive must be a directory")
        self.path.mkdir(parents=True, exist_ok=True, mode=0o700)
        self.retention_days = retention_days
        self.zone = ZoneInfo(timezone_name)
        self._thread_lock = threading.RLock()
        descriptor = os.open(self.path / ".lock", os.O_CREAT | os.O_WRONLY, 0o600)
        os.close(descriptor)

    @contextmanager
    def _locked(self, *, write=False):
        with self._thread_lock, (self.path / ".lock").open("rb") as handle:
            fcntl.flock(handle, fcntl.LOCK_EX if write else fcntl.LOCK_SH)
            try:
                yield
            except pa.ArrowException as error:
                raise StoreError(f"Parquet archive error ({type(error).__name__})") from error
            finally:
                fcntl.flock(handle, fcntl.LOCK_UN)

    def _files(self):
        return sorted(self.path.glob("date=????-??-??/samples.parquet"))

    def _partition(self, stamp):
        date = datetime.fromtimestamp(milliseconds(stamp) / 1000, timezone.utc).date()
        return self.path / f"date={date.isoformat()}" / "samples.parquet"

    @staticmethod
    def _read(path, columns=None):
        # Avoid implicit Hive partition/schema inference from the directory name.
        if not path.exists():
            return []
        parquet = pq.ParquetFile(path, page_checksum_verification=True)
        if not parquet.schema_arrow.equals(SCHEMA, check_metadata=True):
            raise StoreError("Unsupported Parquet archive schema")
        return parquet.read(columns=columns).to_pylist()

    @staticmethod
    def _sync_directory(path):
        descriptor = os.open(path, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(descriptor)
        finally:
            os.close(descriptor)

    def _write(self, path, rows):
        new_directory = not path.parent.exists()
        path.parent.mkdir(parents=True, exist_ok=True, mode=0o700)
        if new_directory:
            self._sync_directory(self.path)
        if not rows:
            path.unlink(missing_ok=True)
            self._sync_directory(path.parent)
            return
        descriptor, temporary = tempfile.mkstemp(prefix=".samples-", suffix=".tmp", dir=path.parent)
        try:
            with os.fdopen(descriptor, "wb") as handle:
                table = pa.Table.from_pylist(sorted(rows, key=lambda row: (row["observed_ms"], row["recorded_ns"], row["id"])), schema=SCHEMA)
                pq.write_table(table, handle, compression="zstd", write_page_checksum=True)
                handle.flush()
                os.fsync(handle.fileno())
            os.replace(temporary, path)
            self._sync_directory(path.parent)
        finally:
            Path(temporary).unlink(missing_ok=True)

    @staticmethod
    def _row(sample, uploaded_at=None):
        coverage = sample_coverage(sample)
        readings = sample.get("readings", {})
        return {
            "id": sample["id"], "source_id": sample["source_id"], "observed_at": sample["observed_at"],
            "observed_ms": coverage[2], "recorded_ns": time.time_ns(), "uploaded_at": uploaded_at,
            "payload": json.dumps(sample, allow_nan=False, separators=(",", ":")),
            **{name: readings.get(name) for name in NUMBERS},
            "boiler_enabled": readings.get("boiler_enabled"), "boiler_reason": readings.get("boiler_reason"),
            "record_kind": "history" if "history_interval" in sample else "live",
            "interval_energy_wh": sample.get("history_interval", {}).get("energy_wh"),
            **{name: bool(value) if name.endswith("_ok") else value
               for name, value in zip(COVERAGE_COLUMNS[1:], coverage[3:])},
        }

    def add(self, sample):
        path, row = self._partition(sample["observed_at"]), self._row(sample)
        with self._locked(write=True):
            rows = self._read(path)
            for existing in rows:
                if existing["id"] == row["id"]:
                    if existing["payload"] != row["payload"]:
                        raise StoreError("Sample ID already contains different data")
                    return
            self._write(path, [*rows, row])

    def pending(self, limit):
        with self._locked():
            rows = []
            for path in self._files():
                index = self._read(path, ["uploaded_at"])
                if any(row["uploaded_at"] is None for row in index):
                    rows.extend(row for row in self._read(path, ["recorded_ns", "id", "payload", "uploaded_at"]) if row["uploaded_at"] is None)
            rows.sort(key=lambda row: (row["recorded_ns"], row["id"]))
            return [json.loads(row["payload"]) for row in rows[:limit]]

    def acknowledge(self, ids, *, now=None):
        wanted, stamp = set(ids), utc_text(now or datetime.now(timezone.utc))
        if not wanted:
            return
        with self._locked(write=True):
            for path in self._files():
                index = self._read(path, ["id", "uploaded_at"])
                if not any(row["id"] in wanted and row["uploaded_at"] is None for row in index):
                    continue
                rows = self._read(path)
                for row in rows:
                    if row["id"] in wanted and row["uploaded_at"] is None:
                        row["uploaded_at"] = stamp
                self._write(path, rows)

    def records(self, start=None, until=None):
        lower, upper = milliseconds(start) if start else None, milliseconds(until) if until else None
        with self._locked():
            rows = []
            for path in self._files():
                rows.extend(row for row in self._read(path, ["observed_ms", "recorded_ns", "id", "payload"])
                            if (lower is None or row["observed_ms"] >= lower) and (upper is None or row["observed_ms"] < upper))
        rows.sort(key=lambda row: (row["observed_ms"], row["recorded_ns"], row["id"]))
        return (json.loads(row["payload"]) for row in rows)

    def replay(self, start=None, until=None):
        lower, upper = milliseconds(start) if start else None, milliseconds(until) if until else None
        count = 0
        with self._locked(write=True):
            for path in self._files():
                rows, changed = self._read(path), False
                for row in rows:
                    if (lower is None or row["observed_ms"] >= lower) and (upper is None or row["observed_ms"] < upper):
                        row["uploaded_at"] = None
                        count += 1
                        changed = True
                if changed:
                    self._write(path, rows)
        return count

    def prune(self, *, now=None):
        """Remove whole local observation days only after all their acknowledgements expire."""
        boundary = retention_cutoff(now or datetime.now(timezone.utc), self.retention_days, self.zone.key)
        cutoff, first_day = utc_text(boundary), boundary.astimezone(self.zone).date()
        deleted = 0
        local_day = lambda row: datetime.fromtimestamp(row["observed_ms"] / 1000, self.zone).date()
        with self._locked(write=True):
            partitions, expired, protected = [], set(), set()
            # A local day can span two UTC partitions. One pending or recently
            # acknowledged record protects the rest of its local day as well.
            for path in self._files():
                index = self._read(path, ["observed_ms", "uploaded_at"])
                days = set()
                for row in index:
                    day = local_day(row)
                    days.add(day)
                    if day < first_day and row["uploaded_at"] is not None and row["uploaded_at"] < cutoff:
                        expired.add(day)
                    else:
                        protected.add(day)
                partitions.append((path, days))
            expired -= protected
            for path, days in partitions:
                if not days & expired:
                    continue
                rows = self._read(path)
                retained = [row for row in rows if local_day(row) not in expired]
                self._write(path, retained)
                deleted += len(rows) - len(retained)
        return deleted

    def status(self):
        with self._locked():
            rows = []
            latest = None
            for path in self._files():
                index = self._read(path, ["observed_at", "observed_ms", "recorded_ns", "uploaded_at"])
                rows.extend(index)
                if index:
                    candidate = max(index, key=lambda row: (row["observed_ms"], row["recorded_ns"]))
                    if latest is None or (candidate["observed_ms"], candidate["recorded_ns"]) > latest[0]:
                        latest = ((candidate["observed_ms"], candidate["recorded_ns"]), path)
            errors = {}
            if latest:
                newest = max(self._read(latest[1], ["observed_ms", "recorded_ns", "payload"]), key=lambda row: (row["observed_ms"], row["recorded_ns"]))
                errors = json.loads(newest["payload"]).get("errors", {})
        return {
            "archive": str(self.path), "format": "parquet", "retention_days": self.retention_days,
            "samples": len(rows), "pending": sum(row["uploaded_at"] is None for row in rows),
            "first_observation": min((row["observed_at"] for row in rows), default=None),
            "last_observation": max((row["observed_at"] for row in rows), default=None),
            "last_upload": max((row["uploaded_at"] for row in rows if row["uploaded_at"]), default=None),
            "last_errors": errors,
        }

    def coverage(self, source_id, timezone_name, days, now=None):
        now = now or datetime.now(timezone.utc)
        first = now.astimezone(ZoneInfo(timezone_name)).date() - timedelta(days=days)
        lower = round(datetime.combine(first, datetime.min.time(), ZoneInfo(timezone_name)).timestamp() * 1000)
        upper = round(now.timestamp() * 1000)
        with self._locked():
            rows = []
            for path in self._files():
                if path.parent.name[5:] < (now - timedelta(days=days + 2)).date().isoformat():
                    continue
                rows.extend(tuple(row[name] for name in COVERAGE_COLUMNS)
                            for row in self._read(path, ["source_id", *COVERAGE_COLUMNS])
                            if row["source_id"] == source_id and lower <= row["observed_ms"] <= upper)
        return daily_coverage(sorted(rows), timezone_name, days, now)

    def last_success(self, source_id):
        with self._locked():
            for path in reversed(self._files()):
                rows = self._read(path, ["source_id", "observed_ms", "power_ok", "boiler_ok"])
                stamps = [row["observed_ms"] for row in rows if row["source_id"] == source_id and row["power_ok"] and row["boiler_ok"]]
                if stamps:
                    return utc_text(datetime.fromtimestamp(max(stamps) / 1000, timezone.utc))
        return None

    def energy_slots(self, source_id, start, until):
        """Return UTC bucket starts already having valid local or recovered energy."""
        lower, upper = milliseconds(start), milliseconds(until)
        with self._locked():
            return {row["observed_ms"] // 300_000 * 300
                    for path in self._files()
                    for row in self._read(path, ["source_id", "observed_ms", "power_ok"])
                    if row["source_id"] == source_id and row["power_ok"] and lower <= row["observed_ms"] < upper}