"""Witness sample storage with local cache plus optional S3/GCS backends."""

from __future__ import annotations

import hashlib
import logging
import os
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Optional, Tuple

logger = logging.getLogger(__name__)


class WitnessStore:
    """Durable witness sample storage for calibration witnesses."""

    def __init__(
        self,
        *,
        backend: str = "local",
        bucket: Optional[str] = None,
        cache_dir: Optional[Path] = None,
        client: Optional[Any] = None,
    ) -> None:
        self.backend = backend.lower()
        self.bucket = bucket
        self.cache_dir = Path(
            cache_dir
            or os.environ.get("WITNESS_CACHE_DIR")
            or "/tmp/loom-witnesses"
        )
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        self.client = client or self._build_client()

        if self.backend in {"s3", "gcs"} and not self.bucket:
            raise ValueError("bucket is required for remote witness storage backends")

    # ------------------------------------------------------------------ #
    def store(self, receipt_id: str, witness_sample: bytes, ts: datetime) -> str:
        """Store witness sample and return SHA-256 hash."""
        ts = self._coerce_timestamp(ts)
        witness_hash = hashlib.sha256(witness_sample).hexdigest()

        cache_path = self._cache_path(receipt_id)
        cache_path.parent.mkdir(parents=True, exist_ok=True)
        cache_path.write_bytes(witness_sample)
        logger.debug("Stored witness in cache: %s", cache_path)

        if self.backend == "s3" and self.client:
            self._put_s3_object(receipt_id, witness_sample, ts)
        elif self.backend == "gcs" and self.client:
            self._put_gcs_blob(receipt_id, witness_sample, ts)

        return witness_hash

    def fetch(self, receipt_id: str, ts: datetime) -> Optional[bytes]:
        """Fetch witness sample (cache-first)."""
        ts = self._coerce_timestamp(ts)
        cache_path = self._cache_path(receipt_id)
        if cache_path.exists():
            logger.debug("Loaded witness from cache: %s", cache_path)
            return cache_path.read_bytes()

        if self.backend == "s3" and self.client:
            data = self._get_s3_object(receipt_id, ts)
            if data is not None:
                cache_path.parent.mkdir(parents=True, exist_ok=True)
                cache_path.write_bytes(data)
                return data

        if self.backend == "gcs" and self.client:
            data = self._get_gcs_blob(receipt_id, ts)
            if data is not None:
                cache_path.parent.mkdir(parents=True, exist_ok=True)
                cache_path.write_bytes(data)
                return data

        return None

    def verify_hash(
        self,
        receipt_id: str,
        ts: datetime,
        expected_hash: str,
    ) -> Tuple[bool, Optional[str]]:
        """Check whether stored witness sample matches expected hash."""
        sample = self.fetch(receipt_id, ts)
        if sample is None:
            return False, f"Witness sample not found for {receipt_id}"

        actual = hashlib.sha256(sample).hexdigest()
        if actual != expected_hash:
            return False, f"Hash mismatch: expected {expected_hash}, got {actual}"
        return True, None

    # ------------------------------------------------------------------ #
    def _build_client(self) -> Optional[Any]:
        if self.backend == "s3":
            try:
                import boto3  # type: ignore

                return boto3.client("s3")
            except Exception as exc:  # pragma: no cover - import guard
                logger.warning("Unable to import boto3 for WitnessStore: %s", exc)
                return None
        if self.backend == "gcs":
            try:
                from google.cloud import storage  # type: ignore

                return storage.Client()
            except Exception as exc:  # pragma: no cover - import guard
                logger.warning(
                    "Unable to import google-cloud-storage for WitnessStore: %s",
                    exc,
                )
                return None
        return None

    def _cache_path(self, receipt_id: str) -> Path:
        return self.cache_dir / receipt_id / "sample.parquet"

    def _object_key(self, receipt_id: str, ts: datetime) -> str:
        return f"{ts.year:04d}/{ts.month:02d}/{ts.day:02d}/{receipt_id}/sample.parquet"

    @staticmethod
    def _coerce_timestamp(ts: datetime) -> datetime:
        if ts.tzinfo is None:
            return ts.replace(tzinfo=timezone.utc)
        return ts.astimezone(timezone.utc)

    # ------------------------------ S3 -------------------------------- #
    def _put_s3_object(self, receipt_id: str, payload: bytes, ts: datetime) -> None:
        key = self._object_key(receipt_id, ts)
        try:
            self.client.put_object(Bucket=self.bucket, Key=key, Body=payload)
            logger.info("Stored witness to s3://%s/%s", self.bucket, key)
        except Exception as exc:  # pragma: no cover - network errors
            logger.error("Failed to write witness to S3: %s", exc)

    def _get_s3_object(self, receipt_id: str, ts: datetime) -> Optional[bytes]:
        key = self._object_key(receipt_id, ts)
        try:
            response = self.client.get_object(Bucket=self.bucket, Key=key)
        except Exception as exc:  # pragma: no cover - network errors
            logger.warning("Failed to fetch witness from S3 (%s): %s", key, exc)
            return None
        body = response.get("Body")
        if body is None:
            logger.warning("S3 object missing Body for %s", key)
            return None
        data = body.read()
        logger.info("Fetched witness from s3://%s/%s", self.bucket, key)
        return data

    # ------------------------------ GCS ------------------------------- #
    def _put_gcs_blob(self, receipt_id: str, payload: bytes, ts: datetime) -> None:
        key = self._object_key(receipt_id, ts)
        try:
            bucket = self.client.bucket(self.bucket)
            blob = bucket.blob(key)
            blob.upload_from_string(payload)
            logger.info("Stored witness to gs://%s/%s", self.bucket, key)
        except Exception as exc:  # pragma: no cover - network errors
            logger.error("Failed to write witness to GCS: %s", exc)

    def _get_gcs_blob(self, receipt_id: str, ts: datetime) -> Optional[bytes]:
        key = self._object_key(receipt_id, ts)
        try:
            bucket = self.client.bucket(self.bucket)
            blob = bucket.blob(key)
            data = blob.download_as_bytes()
            logger.info("Fetched witness from gs://%s/%s", self.bucket, key)
            return data
        except Exception as exc:  # pragma: no cover - network errors
            logger.warning("Failed to fetch witness from GCS (%s): %s", key, exc)
            return None


__all__ = ["WitnessStore"]
