polymarket-insider-tracker/src/polymarket_insider_tracker/storage/repos.py

"""Repository pattern implementations for data access.

This module provides clean data access abstractions for wallet profiles,
funding transfers, and wallet relationships.
"""

from __future__ import annotations

import logging
from dataclasses import dataclass
from datetime import UTC, datetime
from decimal import Decimal
from typing import TYPE_CHECKING

from sqlalchemy import delete, select, update
from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert

from polymarket_insider_tracker.storage.models import (
    FundingTransferModel,
    RiskAssessmentModel,
    WalletProfileModel,
    WalletRelationshipModel,
)

if TYPE_CHECKING:
    from sqlalchemy.ext.asyncio import AsyncSession

logger = logging.getLogger(__name__)


@dataclass
class WalletProfileDTO:
    """Data transfer object for wallet profiles."""

    address: str
    nonce: int
    first_seen_at: datetime | None
    is_fresh: bool
    matic_balance: Decimal | None
    usdc_balance: Decimal | None
    analyzed_at: datetime
    created_at: datetime | None = None
    updated_at: datetime | None = None

    @classmethod
    def from_model(cls, model: WalletProfileModel) -> WalletProfileDTO:
        """Create DTO from SQLAlchemy model."""
        return cls(
            address=model.address,
            nonce=model.nonce,
            first_seen_at=model.first_seen_at,
            is_fresh=model.is_fresh,
            matic_balance=model.matic_balance,
            usdc_balance=model.usdc_balance,
            analyzed_at=model.analyzed_at,
            created_at=model.created_at,
            updated_at=model.updated_at,
        )


@dataclass
class FundingTransferDTO:
    """Data transfer object for funding transfers."""

    from_address: str
    to_address: str
    amount: Decimal
    token: str
    tx_hash: str
    block_number: int
    timestamp: datetime
    created_at: datetime | None = None

    @classmethod
    def from_model(cls, model: FundingTransferModel) -> FundingTransferDTO:
        """Create DTO from SQLAlchemy model."""
        return cls(
            from_address=model.from_address,
            to_address=model.to_address,
            amount=model.amount,
            token=model.token,
            tx_hash=model.tx_hash,
            block_number=model.block_number,
            timestamp=model.timestamp,
            created_at=model.created_at,
        )


@dataclass
class WalletRelationshipDTO:
    """Data transfer object for wallet relationships."""

    wallet_a: str
    wallet_b: str
    relationship_type: str
    confidence: Decimal
    created_at: datetime | None = None

    @classmethod
    def from_model(cls, model: WalletRelationshipModel) -> WalletRelationshipDTO:
        """Create DTO from SQLAlchemy model."""
        return cls(
            wallet_a=model.wallet_a,
            wallet_b=model.wallet_b,
            relationship_type=model.relationship_type,
            confidence=model.confidence,
            created_at=model.created_at,
        )


class WalletRepository:
    """Repository for wallet profile data access.

    Provides CRUD operations for wallet profiles with async support.
    """

    def __init__(self, session: AsyncSession) -> None:
        """Initialize repository with database session.

        Args:
            session: SQLAlchemy async session.
        """
        self.session = session

    async def get_by_address(self, address: str) -> WalletProfileDTO | None:
        """Get wallet profile by address.

        Args:
            address: Wallet address (lowercase).

        Returns:
            WalletProfileDTO if found, None otherwise.
        """
        result = await self.session.execute(
            select(WalletProfileModel).where(WalletProfileModel.address == address.lower())
        )
        model = result.scalar_one_or_none()
        return WalletProfileDTO.from_model(model) if model else None

    async def get_many(self, addresses: list[str]) -> list[WalletProfileDTO]:
        """Get multiple wallet profiles by addresses.

        Args:
            addresses: List of wallet addresses.

        Returns:
            List of WalletProfileDTOs for found addresses.
        """
        normalized = [addr.lower() for addr in addresses]
        result = await self.session.execute(
            select(WalletProfileModel).where(WalletProfileModel.address.in_(normalized))
        )
        return [WalletProfileDTO.from_model(m) for m in result.scalars().all()]

    async def get_fresh_wallets(self, limit: int = 100) -> list[WalletProfileDTO]:
        """Get recent fresh wallets.

        Args:
            limit: Maximum number of results.

        Returns:
            List of WalletProfileDTOs marked as fresh.
        """
        result = await self.session.execute(
            select(WalletProfileModel)
            .where(WalletProfileModel.is_fresh.is_(True))
            .order_by(WalletProfileModel.analyzed_at.desc())
            .limit(limit)
        )
        return [WalletProfileDTO.from_model(m) for m in result.scalars().all()]

    async def upsert(self, dto: WalletProfileDTO) -> WalletProfileDTO:
        """Insert or update wallet profile.

        Args:
            dto: Wallet profile data.

        Returns:
            Updated WalletProfileDTO.
        """
        now = datetime.now(UTC)
        values = {
            "address": dto.address.lower(),
            "nonce": dto.nonce,
            "first_seen_at": dto.first_seen_at,
            "is_fresh": dto.is_fresh,
            "matic_balance": dto.matic_balance,
            "usdc_balance": dto.usdc_balance,
            "analyzed_at": dto.analyzed_at,
            "updated_at": now,
        }

        # Try PostgreSQL upsert first, fall back to SQLite for testing
        try:
            stmt = pg_insert(WalletProfileModel).values(**values, created_at=now)
            stmt = stmt.on_conflict_do_update(
                index_elements=["address"],
                set_={
                    "nonce": stmt.excluded.nonce,
                    "first_seen_at": stmt.excluded.first_seen_at,
                    "is_fresh": stmt.excluded.is_fresh,
                    "matic_balance": stmt.excluded.matic_balance,
                    "usdc_balance": stmt.excluded.usdc_balance,
                    "analyzed_at": stmt.excluded.analyzed_at,
                    "updated_at": stmt.excluded.updated_at,
                },
            )
            await self.session.execute(stmt)
        except Exception:
            # Fall back to SQLite upsert for testing
            sqlite_stmt = sqlite_insert(WalletProfileModel).values(**values, created_at=now)
            sqlite_stmt = sqlite_stmt.on_conflict_do_update(
                index_elements=["address"],
                set_={
                    "nonce": sqlite_stmt.excluded.nonce,
                    "first_seen_at": sqlite_stmt.excluded.first_seen_at,
                    "is_fresh": sqlite_stmt.excluded.is_fresh,
                    "matic_balance": sqlite_stmt.excluded.matic_balance,
                    "usdc_balance": sqlite_stmt.excluded.usdc_balance,
                    "analyzed_at": sqlite_stmt.excluded.analyzed_at,
                    "updated_at": sqlite_stmt.excluded.updated_at,
                },
            )
            await self.session.execute(sqlite_stmt)

        await self.session.flush()
        return dto

    async def delete(self, address: str) -> bool:
        """Delete wallet profile by address.

        Args:
            address: Wallet address.

        Returns:
            True if deleted, False if not found.
        """
        result = await self.session.execute(
            delete(WalletProfileModel).where(WalletProfileModel.address == address.lower())
        )
        # SQLAlchemy Result does have rowcount but typing doesn't reflect it
        return (result.rowcount or 0) > 0  # type: ignore[attr-defined]

    async def mark_stale(self, address: str) -> bool:
        """Mark a wallet profile as stale (soft delete).

        Sets analyzed_at to a very old date to trigger re-analysis.

        Args:
            address: Wallet address.

        Returns:
            True if updated, False if not found.
        """
        stale_time = datetime(2000, 1, 1, tzinfo=UTC)
        result = await self.session.execute(
            update(WalletProfileModel)
            .where(WalletProfileModel.address == address.lower())
            .values(analyzed_at=stale_time, updated_at=datetime.now(UTC))
        )
        # SQLAlchemy Result does have rowcount but typing doesn't reflect it
        return (result.rowcount or 0) > 0  # type: ignore[attr-defined]


class FundingRepository:
    """Repository for funding transfer data access.

    Provides CRUD operations for funding transfers with async support.
    """

    def __init__(self, session: AsyncSession) -> None:
        """Initialize repository with database session.

        Args:
            session: SQLAlchemy async session.
        """
        self.session = session

    async def get_transfers_to(self, address: str, limit: int = 100) -> list[FundingTransferDTO]:
        """Get transfers to a wallet address.

        Args:
            address: Destination wallet address.
            limit: Maximum number of results.

        Returns:
            List of FundingTransferDTOs ordered by timestamp.
        """
        result = await self.session.execute(
            select(FundingTransferModel)
            .where(FundingTransferModel.to_address == address.lower())
            .order_by(FundingTransferModel.timestamp.asc())
            .limit(limit)
        )
        return [FundingTransferDTO.from_model(m) for m in result.scalars().all()]

    async def get_transfers_from(self, address: str, limit: int = 100) -> list[FundingTransferDTO]:
        """Get transfers from a wallet address.

        Args:
            address: Source wallet address.
            limit: Maximum number of results.

        Returns:
            List of FundingTransferDTOs ordered by timestamp.
        """
        result = await self.session.execute(
            select(FundingTransferModel)
            .where(FundingTransferModel.from_address == address.lower())
            .order_by(FundingTransferModel.timestamp.asc())
            .limit(limit)
        )
        return [FundingTransferDTO.from_model(m) for m in result.scalars().all()]

    async def get_first_transfer_to(self, address: str) -> FundingTransferDTO | None:
        """Get the first transfer to a wallet.

        Args:
            address: Wallet address.

        Returns:
            First FundingTransferDTO if found, None otherwise.
        """
        result = await self.session.execute(
            select(FundingTransferModel)
            .where(FundingTransferModel.to_address == address.lower())
            .order_by(FundingTransferModel.timestamp.asc())
            .limit(1)
        )
        model = result.scalar_one_or_none()
        return FundingTransferDTO.from_model(model) if model else None

    async def get_by_tx_hash(self, tx_hash: str) -> FundingTransferDTO | None:
        """Get transfer by transaction hash.

        Args:
            tx_hash: Transaction hash.

        Returns:
            FundingTransferDTO if found, None otherwise.
        """
        result = await self.session.execute(
            select(FundingTransferModel).where(FundingTransferModel.tx_hash == tx_hash.lower())
        )
        model = result.scalar_one_or_none()
        return FundingTransferDTO.from_model(model) if model else None

    async def insert(self, dto: FundingTransferDTO) -> FundingTransferDTO:
        """Insert a new funding transfer.

        Args:
            dto: Funding transfer data.

        Returns:
            Inserted FundingTransferDTO.

        Raises:
            IntegrityError if tx_hash already exists.
        """
        model = FundingTransferModel(
            from_address=dto.from_address.lower(),
            to_address=dto.to_address.lower(),
            amount=dto.amount,
            token=dto.token,
            tx_hash=dto.tx_hash.lower(),
            block_number=dto.block_number,
            timestamp=dto.timestamp,
        )
        self.session.add(model)
        await self.session.flush()
        return dto

    async def insert_many(self, dtos: list[FundingTransferDTO]) -> int:
        """Insert multiple funding transfers.

        Skips duplicates silently.

        Args:
            dtos: List of funding transfer data.

        Returns:
            Number of transfers inserted.
        """
        inserted = 0
        for dto in dtos:
            try:
                await self.insert(dto)
                inserted += 1
            except Exception as e:
                # Skip duplicates
                if "UNIQUE constraint" in str(e) or "duplicate key" in str(e).lower():
                    continue
                raise
        return inserted


class RelationshipRepository:
    """Repository for wallet relationship data access.

    Provides CRUD operations for wallet relationships with async support.
    """

    def __init__(self, session: AsyncSession) -> None:
        """Initialize repository with database session.

        Args:
            session: SQLAlchemy async session.
        """
        self.session = session

    async def get_relationships(
        self, wallet: str, relationship_type: str | None = None
    ) -> list[WalletRelationshipDTO]:
        """Get relationships for a wallet.

        Args:
            wallet: Wallet address.
            relationship_type: Optional filter by type.

        Returns:
            List of WalletRelationshipDTOs.
        """
        stmt = select(WalletRelationshipModel).where(
            (WalletRelationshipModel.wallet_a == wallet.lower())
            | (WalletRelationshipModel.wallet_b == wallet.lower())
        )
        if relationship_type:
            stmt = stmt.where(WalletRelationshipModel.relationship_type == relationship_type)

        result = await self.session.execute(stmt)
        return [WalletRelationshipDTO.from_model(m) for m in result.scalars().all()]

    async def get_related_wallets(
        self, wallet: str, relationship_type: str | None = None
    ) -> list[str]:
        """Get addresses of related wallets.

        Args:
            wallet: Wallet address.
            relationship_type: Optional filter by type.

        Returns:
            List of related wallet addresses.
        """
        relationships = await self.get_relationships(wallet, relationship_type)
        related = set()
        normalized = wallet.lower()
        for rel in relationships:
            if rel.wallet_a == normalized:
                related.add(rel.wallet_b)
            else:
                related.add(rel.wallet_a)
        return list(related)

    async def upsert(self, dto: WalletRelationshipDTO) -> WalletRelationshipDTO:
        """Insert or update wallet relationship.

        Args:
            dto: Wallet relationship data.

        Returns:
            Updated WalletRelationshipDTO.
        """
        now = datetime.now(UTC)
        values = {
            "wallet_a": dto.wallet_a.lower(),
            "wallet_b": dto.wallet_b.lower(),
            "relationship_type": dto.relationship_type,
            "confidence": dto.confidence,
            "created_at": now,
        }

        # Try PostgreSQL upsert first, fall back to SQLite for testing
        try:
            stmt = pg_insert(WalletRelationshipModel).values(**values)
            stmt = stmt.on_conflict_do_update(
                constraint="uq_wallet_relationship",
                set_={"confidence": stmt.excluded.confidence},
            )
            await self.session.execute(stmt)
        except Exception:
            # Fall back to SQLite upsert for testing
            sqlite_stmt = sqlite_insert(WalletRelationshipModel).values(**values)
            sqlite_stmt = sqlite_stmt.on_conflict_do_update(
                index_elements=["wallet_a", "wallet_b", "relationship_type"],
                set_={"confidence": sqlite_stmt.excluded.confidence},
            )
            await self.session.execute(sqlite_stmt)

        await self.session.flush()
        return dto

    async def delete(self, wallet_a: str, wallet_b: str, relationship_type: str) -> bool:
        """Delete a specific relationship.

        Args:
            wallet_a: First wallet address.
            wallet_b: Second wallet address.
            relationship_type: Type of relationship.

        Returns:
            True if deleted, False if not found.
        """
        result = await self.session.execute(
            delete(WalletRelationshipModel).where(
                WalletRelationshipModel.wallet_a == wallet_a.lower(),
                WalletRelationshipModel.wallet_b == wallet_b.lower(),
                WalletRelationshipModel.relationship_type == relationship_type,
            )
        )
        # SQLAlchemy Result does have rowcount but typing doesn't reflect it
        return (result.rowcount or 0) > 0  # type: ignore[attr-defined]


@dataclass
class RiskAssessmentDTO:
    """Data transfer object for a persisted risk assessment.

    Captures everything a future backtest needs without going back to
    public APIs: trade identity, score, per-signal confidences, and
    whether the alert was actually delivered.
    """

    assessment_id: str
    trade_id: str
    wallet_address: str
    market_id: str
    asset_id: str | None
    side: str
    outcome: str | None
    outcome_index: int | None
    price: Decimal
    size: Decimal
    notional_usdc: Decimal
    trade_timestamp: datetime
    weighted_score: Decimal
    signals_triggered: int
    fresh_wallet_confidence: Decimal | None
    size_anomaly_confidence: Decimal | None
    is_niche_market: bool | None
    volume_impact: Decimal | None
    book_impact: Decimal | None
    wallet_age_hours: Decimal | None
    should_alert: bool
    threshold_at_eval: Decimal
    created_at: datetime | None = None


class RiskAssessmentRepository:
    """Repository for risk assessment data access."""

    def __init__(self, session: AsyncSession) -> None:
        self.session = session

    async def insert(self, dto: RiskAssessmentDTO) -> RiskAssessmentDTO:
        """Insert a single assessment. Idempotent on assessment_id collisions."""
        model = RiskAssessmentModel(
            assessment_id=dto.assessment_id,
            trade_id=dto.trade_id,
            wallet_address=dto.wallet_address.lower(),
            market_id=dto.market_id,
            asset_id=dto.asset_id,
            side=dto.side,
            outcome=dto.outcome,
            outcome_index=dto.outcome_index,
            price=dto.price,
            size=dto.size,
            notional_usdc=dto.notional_usdc,
            trade_timestamp=dto.trade_timestamp,
            weighted_score=dto.weighted_score,
            signals_triggered=dto.signals_triggered,
            fresh_wallet_confidence=dto.fresh_wallet_confidence,
            size_anomaly_confidence=dto.size_anomaly_confidence,
            is_niche_market=dto.is_niche_market,
            volume_impact=dto.volume_impact,
            book_impact=dto.book_impact,
            wallet_age_hours=dto.wallet_age_hours,
            should_alert=dto.should_alert,
            threshold_at_eval=dto.threshold_at_eval,
        )
        self.session.add(model)
        await self.session.flush()
        return dto

    async def get_by_assessment_id(self, assessment_id: str) -> RiskAssessmentDTO | None:
        result = await self.session.execute(
            select(RiskAssessmentModel).where(RiskAssessmentModel.assessment_id == assessment_id)
        )
        model = result.scalar_one_or_none()
        if model is None:
            return None
        return RiskAssessmentDTO(
            assessment_id=model.assessment_id,
            trade_id=model.trade_id,
            wallet_address=model.wallet_address,
            market_id=model.market_id,
            asset_id=model.asset_id,
            side=model.side,
            outcome=model.outcome,
            outcome_index=model.outcome_index,
            price=model.price,
            size=model.size,
            notional_usdc=model.notional_usdc,
            trade_timestamp=model.trade_timestamp,
            weighted_score=model.weighted_score,
            signals_triggered=model.signals_triggered,
            fresh_wallet_confidence=model.fresh_wallet_confidence,
            size_anomaly_confidence=model.size_anomaly_confidence,
            is_niche_market=model.is_niche_market,
            volume_impact=model.volume_impact,
            book_impact=model.book_impact,
            wallet_age_hours=model.wallet_age_hours,
            should_alert=model.should_alert,
            threshold_at_eval=model.threshold_at_eval,
            created_at=model.created_at,
        )