Adapters & Repositories¶

class ArcticAdapter(LoggerMixin, DataAdapter):
    """
    ArcticDB storage adapter for OHLCV market data.

    Implements DataAdapter interface for persisting data to ArcticDB on S3.
    Uses dependency injection for testability - accepts optional library client.

    Connection string format: s3s://endpoint:bucket?aws_auth=true

    Attributes:
        SYMBOL_SEPARATOR: Separator used to normalize symbol names

    Example:
        >>> # Default: creates ArcticDB connection lazily
        >>> adapter = ArcticAdapter(
        ...     bucket="my-bucket",
        ...     endpoint="s3.eu-central-1.amazonaws.com",
        ...     library_name="futures"
        ... )

        >>> # For testing: inject mock library
        >>> mock_library = Mock()
        >>> adapter = ArcticAdapter(
        ...     bucket="my-bucket",
        ...     library_name="futures",
        ...     arctic_library=mock_library
        ... )
    """

    SYMBOL_SEPARATOR = "__"

    @staticmethod
    def _get_arcticdb_request_type(type_name: str) -> Any | None:
        """Get request type from arcticdb if installed, else None."""
        try:
            module = importlib.import_module("arcticdb")
        except ImportError:
            return None
        return getattr(module, type_name, None)

    def __init__(
        self,
        bucket: str,
        library_name: str,
        endpoint: str = "s3.eu-central-1.amazonaws.com",
        arctic_library: ArcticLibraryProtocol | None = None,
        *,
        use_ssl: bool = True,
        access_key: str | None = None,
        secret_key: str | None = None,
        region: str = "eu-central-1",
        use_virtual_addressing: bool = True,
    ):
        """
        Initialize ArcticDB adapter.

        Args:
            bucket: S3 bucket name for ArcticDB storage
            library_name: ArcticDB library name (e.g., "futures_1h")
            endpoint: S3 endpoint (default: eu-central-1)
            arctic_library: Optional pre-configured library for testing
            use_ssl: Use SSL/TLS (default: True). Set False for LocalStack.
            access_key: Optional explicit access key (for LocalStack/MinIO)
            secret_key: Optional explicit secret key (for LocalStack/MinIO)
            region: AWS region (default: eu-central-1)
            use_virtual_addressing: Use virtual-hosted style URLs (default: True).
                Set False for LocalStack/MinIO path-style addressing.

        Example:
            >>> # AWS S3 (default - uses aws_auth)
            >>> adapter = ArcticAdapter(bucket="tradai", library_name="futures")

            >>> # LocalStack (explicit credentials, no SSL, path-style)
            >>> adapter = ArcticAdapter(
            ...     bucket="tradai",
            ...     library_name="futures",
            ...     endpoint="localstack:4566",
            ...     use_ssl=False,
            ...     access_key="test",
            ...     secret_key="test",
            ...     use_virtual_addressing=False,
            ... )
        """
        super().__init__()

        # H9: Block LocalStack endpoints in production (SSL enabled = production indicator)
        if use_ssl and endpoint:
            endpoint_lower = endpoint.lower()
            if "localstack" in endpoint_lower or "localhost:4566" in endpoint_lower:
                from tradai.common.exceptions import ConfigurationError

                raise ConfigurationError(
                    f"ArcticDB endpoint '{endpoint}' appears to be LocalStack, "
                    f"but use_ssl=True indicates production. "
                    f"Set use_ssl=False for LocalStack or use a production endpoint."
                )

        self._bucket = bucket
        self._library_name = library_name
        self._endpoint = endpoint
        self._arctic_library: ArcticLibraryProtocol | None = arctic_library
        self._connection: Any = None
        self._use_ssl = use_ssl
        self._access_key = access_key
        self._secret_key = secret_key
        self._region = region
        self._use_virtual_addressing = use_virtual_addressing

    @property
    def library(self) -> ArcticLibraryProtocol:
        """
        Lazy-loaded ArcticDB library.

        Creates connection and gets library on first access.
        Uses aws_auth=true for AWS S3, or explicit credentials for LocalStack/MinIO.

        Returns:
            ArcticDB library instance

        Raises:
            ImportError: If ArcticDB is not installed
        """
        if self._arctic_library is None:
            try:
                from arcticdb import Arctic
            except ImportError as e:
                raise ImportError(
                    "ArcticDB is not installed. Install with: pip install arcticdb "
                    "(only available on Linux x86_64 and Windows)"
                ) from e

            s3_url = self._build_connection_string()
            self._connection = Arctic(s3_url)
            self._arctic_library = self._connection.get_library(
                self._library_name,
                create_if_missing=True,
            )
            self.logger.info(f"Connected to ArcticDB: {self._bucket}/{self._library_name}")

        return self._arctic_library

    def _build_connection_string(self) -> str:
        """
        Build ArcticDB S3 connection string.

        Format: s3[s]://host:bucket?port=X&params

        For LocalStack/MinIO, endpoint can be "host:port" and the port
        will be extracted and added as a query parameter.

        Returns:
            Connection string for ArcticDB
        """
        protocol = "s3s" if self._use_ssl else "s3"

        # Parse endpoint to extract host and optional port
        host = self._endpoint
        port: str | None = None
        if ":" in self._endpoint:
            parts = self._endpoint.rsplit(":", 1)
            # Check if the part after : is a port number (not part of hostname)
            if parts[1].isdigit():
                host = parts[0]
                port = parts[1]

        # Build query parameters
        params: list[str] = []

        # Add port if specified
        if port:
            params.append(f"port={port}")

        if self._access_key and self._secret_key:
            # Explicit credentials (LocalStack/MinIO)
            params.append(f"access={self._access_key}")
            params.append(f"secret={self._secret_key}")
            params.append(f"region={self._region}")
        else:
            # AWS credential delegation
            params.append("aws_auth=true")

        # Add virtual addressing option (false for LocalStack/MinIO)
        if not self._use_virtual_addressing:
            params.append("use_virtual_addressing=false")

        query_string = "&".join(params)
        return f"{protocol}://{host}:{self._bucket}?{query_string}"

    def save(
        self,
        data: OHLCVData,
        symbols: SymbolList,
        latest_query_date: datetime,
        timeframe: str | None = None,
    ) -> None:
        """
        Save OHLCV data to ArcticDB.

        Groups data by symbol and uses upsert semantics:
        - New symbols: uses write
        - Existing symbols: uses update (upserts rows by index)

        Stores metadata_version, last_query_date, last_candle_date,
        and optional timeframe in metadata for incremental updates.

        Args:
            data: Market data to save
            symbols: Symbols being saved (for metadata)
            latest_query_date: Timestamp of when data was queried
            timeframe: Optional CCXT timeframe string stored in metadata

        Raises:
            StorageError: If saving fails

        Example:
            >>> adapter.save(
            ...     data=ohlcv_data,
            ...     symbols=SymbolList.from_input(["BTC/USDT:USDT"]),
            ...     latest_query_date=datetime.now(UTC),
            ...     timeframe="1h",
            ... )
        """
        if data.is_empty:
            self.logger.warning("Attempted to save empty OHLCVData, skipping")
            return

        try:
            df = data.to_dataframe()
            grouped = df.groupby("symbol")

            saved_count = 0
            for symbol in symbols.to_list():
                if symbol not in grouped.groups:
                    continue

                symbol_df = grouped.get_group(symbol).copy()

                # Prepare for ArcticDB: drop symbol column, set date as index
                prepared_df = symbol_df.drop(columns=["symbol"]).set_index("date")

                arctic_symbol = self._normalize_symbol(symbol)
                metadata = self._build_symbol_metadata(
                    latest_query_date=latest_query_date,
                    last_candle_date=prepared_df.index.max(),
                    timeframe=timeframe,
                )

                # Use upsert semantics: update existing, write new
                if self.library.has_symbol(arctic_symbol):
                    # Update existing symbol (upserts rows by index)
                    self.library.update(
                        arctic_symbol,
                        prepared_df,
                        metadata=metadata,
                        upsert=True,
                        prune_previous_versions=True,
                    )
                    self.logger.debug(f"Updated symbol: {symbol}")
                else:
                    # Write new symbol
                    self.library.write(
                        arctic_symbol,
                        prepared_df,
                        metadata=metadata,
                        prune_previous_versions=True,
                    )
                    self.logger.debug(f"Created symbol: {symbol}")

                saved_count += 1

            if saved_count > 0:
                self.logger.info(f"Saved {saved_count} symbols, {data.row_count} rows")

        except (KeyError, ValueError, AttributeError) as e:
            raise StorageError(f"Data format error during save: {e}") from e
        except Exception as e:
            self.logger.error("Unexpected error during save: %s", e)
            raise StorageError(f"Failed to save data: {e}") from e

    def save_batch(
        self,
        data: OHLCVData,
        symbols: SymbolList,
        latest_query_date: datetime,
        timeframe: str | None = None,
    ) -> int:
        """
        Save OHLCV data to ArcticDB using batch write for better performance.

        Uses write_batch for 2-3x faster writes compared to sequential save().
        Best for initial data loads or when updating multiple symbols.

        Note: This method uses write (not update), so it replaces existing data.
        For incremental upserts, use save() instead.

        Args:
            data: Market data to save
            symbols: Symbols being saved (for metadata)
            latest_query_date: Timestamp of when data was queried
            timeframe: Optional CCXT timeframe string stored in metadata

        Returns:
            Number of symbols successfully saved

        Raises:
            StorageError: If batch save fails

        Example:
            >>> # Save 10 symbols in batch (2-3x faster than save())
            >>> count = adapter.save_batch(
            ...     data=ohlcv_data,
            ...     symbols=SymbolList.from_input(["BTC/USDT:USDT", ...]),
            ...     latest_query_date=datetime.now(UTC),
            ...     timeframe="1h",
            ... )
            >>> print(f"Saved {count} symbols")
        """
        if data.is_empty:
            self.logger.warning("Attempted to save_batch empty OHLCVData, skipping")
            return 0

        try:
            payloads = self._prepare_write_payloads(data, symbols, latest_query_date, timeframe)

            if not payloads:
                self.logger.warning("No data to save in batch")
                return 0

            success_count = self._execute_batch_write(payloads, data.row_count)
            return success_count

        except (KeyError, ValueError, AttributeError) as e:
            raise StorageError(f"Data format error during batch save: {e}") from e
        except Exception as e:
            self.logger.error("Unexpected error during batch save: %s", e)
            raise StorageError(f"Failed to batch save data: {e}") from e

    def _prepare_write_payloads(
        self,
        data: OHLCVData,
        symbols: SymbolList,
        latest_query_date: datetime,
        timeframe: str | None = None,
    ) -> list[Any]:
        """
        Prepare WritePayload objects for batch write operation.

        Args:
            data: Market data to prepare
            symbols: Symbols to prepare payloads for
            latest_query_date: Timestamp of when data was queried
            timeframe: Optional CCXT timeframe string stored in metadata

        Returns:
            List of WritePayload objects ready for batch write
        """
        ArcticWritePayload = self._get_arcticdb_request_type("WritePayload") or WritePayload

        df = data.to_dataframe()
        grouped = df.groupby("symbol")

        payloads = []
        for symbol in symbols.to_list():
            if symbol not in grouped.groups:
                continue

            symbol_df = grouped.get_group(symbol).copy()

            # Prepare for ArcticDB: drop symbol column, set date as index
            prepared_df = symbol_df.drop(columns=["symbol"]).set_index("date")

            arctic_symbol = self._normalize_symbol(symbol)
            metadata = self._build_symbol_metadata(
                latest_query_date=latest_query_date,
                last_candle_date=prepared_df.index.max(),
                timeframe=timeframe,
            )

            payloads.append(
                ArcticWritePayload(symbol=arctic_symbol, data=prepared_df, metadata=metadata)
            )

        return payloads

    def _execute_batch_write(self, payloads: list[Any], row_count: int) -> int:
        """
        Execute batch write and count successful writes.

        Args:
            payloads: List of WritePayload objects to write
            row_count: Total number of rows being written (for logging)

        Returns:
            Number of symbols successfully written
        """
        results = self.library.write_batch(payloads)

        # Count successes and log errors
        success_count = 0
        for result in results:
            # Check for DataError (has error_code attribute)
            if hasattr(result, "error_code"):
                self.logger.error(f"Batch write error for {result.symbol}: {result}")
            else:
                success_count += 1

        self.logger.info(f"Batch saved {success_count}/{len(payloads)} symbols, {row_count} rows")
        return success_count

    def load(
        self,
        symbols: SymbolList,
        date_range: DateRange,
    ) -> OHLCVData:
        """
        Load OHLCV data from ArcticDB.

        Uses read_batch for efficient multi-symbol loading.
        Reconstructs DataFrame with symbol column.

        Args:
            symbols: Symbols to load
            date_range: Date range to load

        Returns:
            OHLCVData from storage

        Raises:
            DataNotFoundError: If no data found for symbols
            StorageError: If loading fails

        Example:
            >>> data = adapter.load(
            ...     symbols=SymbolList.from_input(["BTC/USDT:USDT"]),
            ...     date_range=DateRange.from_strings("2024-01-01", "2024-01-31")
            ... )
        """
        try:
            # Prefer real arcticdb request types when available (production),
            # but fall back to local models for macOS/CI environments where
            # arcticdb isn't installable (tests inject a mock library).
            ArcticReadRequest = self._get_arcticdb_request_type("ReadRequest") or ReadRequest

            read_requests = [
                ArcticReadRequest(
                    symbol=self._normalize_symbol(symbol),
                    date_range=(date_range.start, date_range.end),
                )
                for symbol in symbols.to_list()
            ]

            results = self.library.read_batch(read_requests)

            dataframes = []
            for result in results:
                # Skip DataError results (check for 'data' attribute)
                if not hasattr(result, "data") or result.data is None:
                    continue
                if hasattr(result, "error_code"):
                    # This is a DataError, skip it
                    continue

                df = result.data.reset_index()  # date back as column

                # Denormalize symbol name back to trading format
                trading_symbol = self._denormalize_symbol(result.symbol)
                df.insert(0, "symbol", trading_symbol)

                dataframes.append(df)

            if not dataframes:
                raise DataNotFoundError(
                    f"No data found for {symbols.to_list()} "
                    f"between {date_range.start} and {date_range.end}"
                )

            combined_df = pd.concat(dataframes, ignore_index=True)
            return OHLCVData.from_dataframe(combined_df)

        except DataNotFoundError:
            raise
        except (KeyError, ValueError, AttributeError) as e:
            raise StorageError(f"Data format error during load: {e}") from e
        except Exception as e:
            self.logger.error("Unexpected error during load: %s", e)
            raise StorageError(f"Failed to load data: {e}") from e

    def exists(self, symbols: SymbolList) -> dict[str, bool]:
        """
        Check if data exists for given symbols.

        Uses has_symbol for efficient individual checks.

        Args:
            symbols: Symbols to check

        Returns:
            Dictionary mapping symbol -> exists (True/False)

        Example:
            >>> result = adapter.exists(
            ...     SymbolList.from_input(["BTC/USDT:USDT", "ETH/USDT:USDT"])
            ... )
            >>> result
            {'BTC/USDT:USDT': True, 'ETH/USDT:USDT': False}
        """
        return {
            symbol: self.library.has_symbol(self._normalize_symbol(symbol))
            for symbol in symbols.to_list()
        }

    def get_latest_date(self, symbols: SymbolList) -> dict[str, datetime]:
        """
        Get latest available date for each symbol.

        Returns the actual last candle date from stored data.
        Falls back to last_query_date for backwards compatibility with older data.

        Args:
            symbols: Symbols to check

        Returns:
            Dictionary mapping symbol -> latest candle date (only for existing symbols)

        Example:
            >>> result = adapter.get_latest_date(
            ...     SymbolList.from_input(["BTC/USDT:USDT"])
            ... )
            >>> result
            {'BTC/USDT:USDT': datetime(2024, 1, 31, 23, 0, 0)}
        """
        # Prefer real arcticdb request types when available (production),
        # but fall back to local models for macOS/CI environments.
        ArcticReadInfoRequest = (
            self._get_arcticdb_request_type("ReadInfoRequest") or ReadInfoRequest
        )

        read_info_requests = [
            ArcticReadInfoRequest(symbol=self._normalize_symbol(symbol))
            for symbol in symbols.to_list()
        ]

        results = self.library.read_metadata_batch(read_info_requests)

        latest_dates: dict[str, datetime] = {}
        for result in results:
            # Skip DataError results
            if hasattr(result, "error_code"):
                continue
            if not hasattr(result, "metadata") or result.metadata is None:
                continue

            trading_symbol = self._denormalize_symbol(result.symbol)

            if "metadata_version" not in result.metadata:
                self.logger.debug(f"Symbol '{trading_symbol}' has pre-versioning metadata")

            # Prefer last_candle_date (actual data end date)
            # Fall back to last_query_date for backwards compatibility
            date_str = result.metadata.get("last_candle_date") or result.metadata.get(
                "last_query_date"
            )
            if date_str:
                latest_dates[trading_symbol] = ensure_utc_naive(datetime.fromisoformat(date_str))

        return latest_dates

    def get_stored_timeframe(self, symbols: SymbolList) -> dict[str, str | None]:
        """Get stored timeframe for each symbol from ArcticDB metadata.

        Returns None for legacy data without timeframe metadata.

        Args:
            symbols: Symbols to check

        Returns:
            Dictionary mapping symbol -> timeframe string or None
        """
        ArcticReadInfoRequest = (
            self._get_arcticdb_request_type("ReadInfoRequest") or ReadInfoRequest
        )

        read_info_requests = [
            ArcticReadInfoRequest(symbol=self._normalize_symbol(symbol))
            for symbol in symbols.to_list()
        ]

        results = self.library.read_metadata_batch(read_info_requests)

        timeframes: dict[str, str | None] = {}
        for result in results:
            if hasattr(result, "error_code"):
                continue
            if not hasattr(result, "metadata") or result.metadata is None:
                continue

            trading_symbol = self._denormalize_symbol(result.symbol)
            timeframes[trading_symbol] = result.metadata.get("timeframe")

        return timeframes

    def list_symbols(self) -> list[str]:
        """
        List all symbols stored in ArcticDB.

        Returns symbols in trading format (e.g., "BTC/USDT:USDT").

        Returns:
            List of trading symbols stored in the library

        Example:
            >>> symbols = adapter.list_symbols()
            >>> symbols
            ['BTC/USDT:USDT', 'ETH/USDT:USDT']
        """
        arctic_symbols = self.library.list_symbols()
        return [self._denormalize_symbol(s) for s in arctic_symbols]

    def delete(self, symbols: SymbolList) -> dict[str, bool]:
        """
        Delete data for given symbols from ArcticDB.

        Args:
            symbols: Symbols to delete

        Returns:
            Dictionary mapping symbol -> deleted (True if existed and deleted)

        Raises:
            StorageError: If deletion fails

        Example:
            >>> result = adapter.delete(
            ...     SymbolList.from_input(["BTC/USDT:USDT"])
            ... )
            >>> result
            {'BTC/USDT:USDT': True}
        """
        try:
            results: dict[str, bool] = {}
            for symbol in symbols.to_list():
                arctic_symbol = self._normalize_symbol(symbol)
                if self.library.has_symbol(arctic_symbol):
                    self.library.delete(arctic_symbol)
                    results[symbol] = True
                    self.logger.info(f"Deleted symbol: {symbol}")
                else:
                    results[symbol] = False
            return results
        except Exception as e:
            raise StorageError(f"Failed to delete symbols: {e}") from e

    @staticmethod
    def _build_symbol_metadata(
        latest_query_date: datetime,
        last_candle_date: datetime,
        timeframe: str | None = None,
    ) -> dict[str, Any]:
        """Build metadata dict for a single symbol write.

        Args:
            latest_query_date: Timestamp of when data was queried
            last_candle_date: Timestamp of the last candle in the data
            timeframe: Optional CCXT timeframe string (e.g., "1h")

        Returns:
            Metadata dict with version, query date, candle date, and optional timeframe
        """
        meta: dict[str, Any] = {
            "metadata_version": 2,
            "last_query_date": latest_query_date.isoformat(),
            "last_candle_date": last_candle_date.isoformat(),
        }
        if timeframe is not None:
            meta["timeframe"] = timeframe
        return meta

    @classmethod
    def _normalize_symbol(cls, symbol: str) -> str:
        """
        Convert trading symbol to ArcticDB-safe format.

        ArcticDB symbols cannot contain / or : characters.

        Args:
            symbol: Trading symbol (e.g., "BTC/USDT:USDT")

        Returns:
            Normalized symbol (e.g., "BTC__USDT__USDT")

        Example:
            >>> ArcticAdapter._normalize_symbol("BTC/USDT:USDT")
            'BTC__USDT__USDT'
        """
        return symbol.replace("/", cls.SYMBOL_SEPARATOR).replace(":", cls.SYMBOL_SEPARATOR)

    @classmethod
    def _denormalize_symbol(cls, arctic_symbol: str) -> str:
        """
        Convert ArcticDB symbol back to trading format.

        Args:
            arctic_symbol: ArcticDB symbol (e.g., "BTC__USDT__USDT")

        Returns:
            Trading symbol (e.g., "BTC/USDT:USDT")

        Example:
            >>> ArcticAdapter._denormalize_symbol("BTC__USDT__USDT")
            'BTC/USDT:USDT'
        """
        parts = arctic_symbol.split(cls.SYMBOL_SEPARATOR)
        if len(parts) >= 3:
            # Futures: BTC__USDT__USDT -> BTC/USDT:USDT
            settle = cls.SYMBOL_SEPARATOR.join(parts[2:])
            return f"{parts[0]}/{parts[1]}:{settle}"
        elif len(parts) == 2:
            # Spot: BTC__USDT -> BTC/USDT
            return f"{parts[0]}/{parts[1]}"
        return arctic_symbol