"""GCS request logging for EB1-GNN server.

This module provides async-safe buffered logging to Google Cloud Storage.
Logs are buffered in memory and flushed to GCS either:
- Every 60 seconds (time-based)
- Every 100 entries (count-based)
- Whichever comes first

Path structure: gs://{bucket}/requests/{YYYY-MM-DD}/requests_{HH}.jsonl

The logger is designed for high-throughput production use:
- Non-blocking: log() returns immediately
- Batched writes: reduces GCS API calls
- Graceful shutdown: flushes remaining logs on close
"""

from __future__ import annotations

import atexit
import json
import logging
import threading
import time
from collections import deque
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

# Default configuration
DEFAULT_BUCKET = "eb1-logs"
DEFAULT_FLUSH_INTERVAL_SECONDS = 60
DEFAULT_FLUSH_COUNT = 100


class GCSRequestLogger:
    """Buffered logger that writes request analytics to GCS.

    Thread-safe implementation that buffers log entries and periodically
    flushes them to GCS in JSONL format.

    Example:
        gcs_logger = GCSRequestLogger(bucket="eb1-logs")
        gcs_logger.start()

        # Log requests
        gcs_logger.log({"request_id": "abc", "model": "eb1-preview", ...})

        # On shutdown
        gcs_logger.stop()
    """

    def __init__(
        self,
        bucket: str = DEFAULT_BUCKET,
        flush_interval_seconds: int = DEFAULT_FLUSH_INTERVAL_SECONDS,
        flush_count: int = DEFAULT_FLUSH_COUNT,
        enabled: bool = True,
    ) -> None:
        """Initialize the GCS logger.

        Args:
            bucket: GCS bucket name (without gs:// prefix).
            flush_interval_seconds: Maximum time between flushes.
            flush_count: Maximum entries before triggering flush.
            enabled: If False, logs are discarded (for local dev).
        """
        self._bucket_name = bucket
        self._flush_interval = flush_interval_seconds
        self._flush_count = flush_count
        self._enabled = enabled

        self._buffer: deque[dict[str, Any]] = deque()
        self._lock = threading.Lock()
        self._stop_event = threading.Event()
        self._flush_thread: threading.Thread | None = None
        self._client: Any = None  # Lazy-loaded GCS client

        # Register shutdown handler
        atexit.register(self.stop)

    def start(self) -> None:
        """Start the background flush thread."""
        if not self._enabled:
            logger.info("GCS logging disabled")
            return

        if self._flush_thread is not None:
            return

        self._stop_event.clear()
        self._flush_thread = threading.Thread(
            target=self._flush_loop,
            name="gcs-logger-flush",
            daemon=True,
        )
        self._flush_thread.start()
        logger.info(
            "GCS logger started: bucket=%s, flush_interval=%ds, flush_count=%d",
            self._bucket_name,
            self._flush_interval,
            self._flush_count,
        )

    def stop(self) -> None:
        """Stop the flush thread and flush remaining logs."""
        if self._flush_thread is None:
            return

        self._stop_event.set()
        self._flush_thread.join(timeout=10)
        self._flush_thread = None

        # Final flush
        self._flush()
        logger.info("GCS logger stopped")

    def log(self, entry: dict[str, Any]) -> None:
        """Add a log entry to the buffer.

        This method is thread-safe and non-blocking.

        Args:
            entry: Dictionary containing request analytics.
        """
        if not self._enabled:
            return

        with self._lock:
            self._buffer.append(entry)
            buffer_size = len(self._buffer)

        # Trigger immediate flush if count threshold reached
        if buffer_size >= self._flush_count:
            self._flush()

    def _flush_loop(self) -> None:
        """Background thread that periodically flushes the buffer."""
        while not self._stop_event.is_set():
            # Wait for flush interval or stop signal
            self._stop_event.wait(timeout=self._flush_interval)
            if not self._stop_event.is_set():
                self._flush()

    def _flush(self) -> None:
        """Flush buffered entries to GCS."""
        # Atomically swap buffer
        with self._lock:
            if not self._buffer:
                return
            entries = list(self._buffer)
            self._buffer.clear()

        if not entries:
            return

        try:
            self._write_to_gcs(entries)
            logger.debug("Flushed %d entries to GCS", len(entries))
        except Exception as e:
            logger.error("Failed to flush to GCS: %s", e)
            # Put entries back in buffer for retry
            with self._lock:
                for entry in reversed(entries):
                    self._buffer.appendleft(entry)

    def _write_to_gcs(self, entries: list[dict[str, Any]]) -> None:
        """Write entries to GCS as JSONL.

        Path: gs://{bucket}/requests/{YYYY-MM-DD}/requests_{HH}.jsonl

        Args:
            entries: List of log entries to write.
        """
        # Lazy-load GCS client
        if self._client is None:
            try:
                from google.cloud import storage

                self._client = storage.Client()
            except ImportError:
                logger.warning(
                    "google-cloud-storage not installed. "
                    "Install with: pip install google-cloud-storage"
                )
                return
            except Exception as e:
                logger.error("Failed to initialize GCS client: %s", e)
                return

        # Build path: requests/2026-01-05/requests_14.jsonl
        now = datetime.now(timezone.utc)
        date_str = now.strftime("%Y-%m-%d")
        hour_str = now.strftime("%H")
        blob_path = f"requests/{date_str}/requests_{hour_str}.jsonl"

        # Convert entries to JSONL
        jsonl_content = "\n".join(json.dumps(entry) for entry in entries) + "\n"

        try:
            bucket = self._client.bucket(self._bucket_name)
            blob = bucket.blob(blob_path)

            # Append to existing blob or create new
            # Note: GCS doesn't support true append, so we download + append + upload
            # For production, consider using Cloud Logging or BigQuery streaming
            existing_content = ""
            if blob.exists():
                existing_content = blob.download_as_text()

            blob.upload_from_string(
                existing_content + jsonl_content,
                content_type="application/x-ndjson",
            )

            logger.debug(
                "Wrote %d entries to gs://%s/%s",
                len(entries),
                self._bucket_name,
                blob_path,
            )

        except Exception as e:
            logger.error("GCS write failed: %s", e)
            raise

    @property
    def buffer_size(self) -> int:
        """Return current buffer size."""
        with self._lock:
            return len(self._buffer)


# Global logger instance (initialized by server)
_global_logger: GCSRequestLogger | None = None


def get_gcs_logger() -> GCSRequestLogger | None:
    """Get the global GCS logger instance."""
    return _global_logger


def init_gcs_logger(
    bucket: str = DEFAULT_BUCKET,
    enabled: bool = True,
    flush_interval: int = DEFAULT_FLUSH_INTERVAL_SECONDS,
    flush_count: int = DEFAULT_FLUSH_COUNT,
) -> GCSRequestLogger:
    """Initialize and start the global GCS logger.

    Args:
        bucket: GCS bucket name.
        enabled: Whether logging is enabled.
        flush_interval: Seconds between flushes.
        flush_count: Max entries before flush.

    Returns:
        The initialized GCS logger.
    """
    global _global_logger

    if _global_logger is not None:
        _global_logger.stop()

    _global_logger = GCSRequestLogger(
        bucket=bucket,
        enabled=enabled,
        flush_interval_seconds=flush_interval,
        flush_count=flush_count,
    )
    _global_logger.start()
    return _global_logger


def log_request(entry: dict[str, Any]) -> None:
    """Log a request to GCS (convenience function).

    Args:
        entry: Request analytics dictionary.
    """
    if _global_logger is not None:
        _global_logger.log(entry)
