csv_token_input_source.py

"""
CSV input source for processing historical token launch data
Reads from CSV file in batches and generates TokenLaunchEvent instances
Deletes processed batches from CSV file for progress tracking
"""

import csv
import logging
from datetime import datetime
from pathlib import Path
from typing import Dict

from src.domain.input_source.base_input_source import InputSource
from src.domain.model.events.token_launch_event import TokenLaunchEvent

logger = logging.getLogger(__name__)


class CSVTokenInputSource(InputSource):
    """Input source that reads CSV file in batches and generates TokenLaunchEvent instances with batch deletion"""

    def __init__(self, csv_file: str = "res/token_dataset.csv", batch_size: int = 100):
        """
        Initialize the CSV token input source

        Args:
            csv_file: Path to CSV file containing token data (default: res/token_dataset.csv)
            batch_size: Number of rows to read and process at a time (default: 100)
        """
        from src.constants import PROJECT_ROOT

        csv_path = Path(csv_file)
        if csv_path.is_absolute():
            self.csv_file = csv_path
        else:
            self.csv_file = PROJECT_ROOT / csv_file

        self.batch_size = batch_size
        self.is_running = False
        self.total_rows_initial = 0
        self.processed_rows = 0
        self.failed_rows = 0
        self.fieldnames = []
        self.current_batch = []
        self.current_batch_index = 0

    async def initialize(self, callback=None):
        """Initialize the input source with callback"""
        await super().initialize(callback)

        # Just verify file exists and get total count
        if self.csv_file.exists():
            try:
                # Count total rows without loading into memory
                with open(self.csv_file, "r", encoding="utf-8") as f:
                    # Read header to get fieldnames
                    reader = csv.DictReader(f)
                    self.fieldnames = reader.fieldnames

                    # Count rows
                    row_count = sum(1 for _ in reader)
                    self.total_rows_initial = row_count

                logger.info(f"📁 CSV file found: {self.csv_file}")
                logger.info(f"📊 Total rows to process: {self.total_rows_initial}")
                logger.info(f"📦 Batch size: {self.batch_size}")

                # Log column names for debugging
                logger.debug(f"📋 CSV columns: {self.fieldnames}")

                # Verify essential columns exist
                essential_columns = ["mint", "name", "symbol"]
                missing_columns = [
                    col for col in essential_columns if col not in self.fieldnames
                ]
                if missing_columns:
                    logger.warning(f"⚠️ Missing essential columns: {missing_columns}")
                    logger.warning(f"   Available columns: {self.fieldnames}")

            except Exception as e:
                logger.error(f"❌ Failed to initialize CSV file: {e}")
                self.total_rows_initial = 0
        else:
            logger.error(f"❌ CSV file not found: {self.csv_file}")
            self.total_rows_initial = 0

    def _read_next_batch(self):
        """Read the next batch of rows from CSV file"""
        try:
            with open(self.csv_file, "r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                self.fieldnames = reader.fieldnames

                # Read up to batch_size rows
                batch = []
                for i, row in enumerate(reader):
                    if i >= self.batch_size:
                        break
                    batch.append(row)

                return batch

        except Exception as e:
            logger.error(f"❌ Failed to read batch from CSV: {e}")
            return []

    def _delete_processed_batch(self, rows_to_delete: int):
        """Delete the first rows_to_delete rows from CSV file"""
        try:
            # Read all remaining rows (after the batch we're deleting)
            remaining_rows = []
            with open(self.csv_file, "r", encoding="utf-8") as f:
                reader = csv.DictReader(f)
                fieldnames = reader.fieldnames

                # Skip the first rows_to_delete rows
                for i, row in enumerate(reader):
                    if i >= rows_to_delete:
                        remaining_rows.append(row)

            # Write back the remaining rows
            with open(self.csv_file, "w", newline="", encoding="utf-8") as f:
                if remaining_rows or fieldnames:
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    writer.writeheader()
                    writer.writerows(remaining_rows)
                else:
                    # If no rows left, just write header
                    writer = csv.DictWriter(f, fieldnames=fieldnames)
                    writer.writeheader()

            logger.info(
                f"📝 Deleted {rows_to_delete} processed rows from CSV. {len(remaining_rows)} rows remaining"
            )

        except Exception as e:
            logger.error(f"❌ Failed to delete processed batch from CSV: {e}")

    async def start(self):
        """Start processing CSV rows in batches"""
        if self.total_rows_initial == 0:
            logger.error(f"❌ No data to process from CSV file")
            return

        self.is_running = True
        logger.info(f"🚀 Starting CSV batch processing from {self.csv_file}")
        logger.info(
            f"📊 Processing {self.total_rows_initial} tokens in batches of {self.batch_size}"
        )

        while self.is_running:
            # Read next batch
            batch = self._read_next_batch()

            if not batch:
                # No more data
                logger.info(f"✅ All batches processed")
                break

            logger.info(f"📦 Processing batch of {len(batch)} tokens...")

            # Process each row in the batch
            batch_processed = 0
            batch_failed = 0
            rows_to_delete = 0  # Track how many rows we actually processed

            for row in batch:
                if not self.is_running:
                    # CRITICAL: Don't process more rows if stopped, but don't lose the ones we haven't processed
                    logger.warning(
                        f"⚠️ Processing stopped mid-batch. Processed {rows_to_delete}/{len(batch)} rows in this batch."
                    )
                    break

                try:
                    # Create TokenLaunchEvent
                    event = self._create_token_launch_event(row)

                    # Emit event through callback
                    if self._on_new_event:
                        try:
                            await self._on_new_event(event)
                            batch_processed += 1
                            self.processed_rows += 1
                            rows_to_delete += (
                                1  # Successfully processed, safe to delete
                            )

                            # Log individual progress
                            progress_pct = (
                                self.processed_rows / self.total_rows_initial
                            ) * 100
                            logger.debug(
                                f"  ✅ {row.get('name', 'Unknown')} ({row.get('symbol', '???')}). "
                                f"Total progress: {self.processed_rows}/{self.total_rows_initial} ({progress_pct:.1f}%)"
                            )

                        except Exception as e:
                            logger.error(f"  ❌ Error processing row: {e}")
                            logger.error(
                                f"     Row: {row.get('name', 'Unknown')} ({row.get('symbol', '???')})"
                            )
                            batch_failed += 1
                            self.failed_rows += 1
                            rows_to_delete += 1  # Failed but handled, safe to delete

                            # Save failed row to separate file
                            self._save_failed_row(row, str(e))

                except Exception as e:
                    logger.error(f"  ❌ Failed to create event from row: {e}")
                    batch_failed += 1
                    self.failed_rows += 1
                    rows_to_delete += 1  # Failed but handled, safe to delete

            # CRITICAL FIX: Only delete the rows we actually processed
            if rows_to_delete > 0:
                self._delete_processed_batch(rows_to_delete)

                logger.info(
                    f"📦 Batch complete: {batch_processed} successful, {batch_failed} failed, {rows_to_delete} deleted. "
                    f"Overall: {self.processed_rows}/{self.total_rows_initial} processed"
                )
            else:
                logger.warning(
                    f"⚠️ No rows processed in this batch - no deletion performed"
                )

        if self.is_running:
            logger.info(
                f"✅ CSV processing complete: "
                f"{self.processed_rows}/{self.total_rows_initial} successful, "
                f"{self.failed_rows} failed"
            )
        else:
            logger.info(
                f"🛑 CSV processing stopped: "
                f"{self.processed_rows}/{self.total_rows_initial} successful, "
                f"{self.failed_rows} failed"
            )

    def _create_token_launch_event(self, row: Dict[str, str]) -> TokenLaunchEvent:
        """
        Create TokenLaunchEvent from CSV row

        Maps CSV columns to the format expected by TokenMetadataStage:
        - Essential: mint, name, symbol
        - Optional: description, metadata_uri, image_uri
        - All other columns preserved in raw_data
        """
        # Row is already a dict
        row_dict = row

        # Create raw_data with proper field mapping
        raw_data = {
            # Essential fields (TokenMetadataStage looks for these)
            "mint": row_dict.get("mint", ""),
            "name": row_dict.get("name", ""),
            "symbol": row_dict.get("symbol", ""),
            # Optional fields with proper key names
            "description": row_dict.get("description", ""),
            "uri": row_dict.get(
                "metadata_uri", ""
            ),  # TokenMetadataStage looks for 'uri'
            "metadata_uri": row_dict.get(
                "metadata_uri", ""
            ),  # Also include as metadata_uri
            "image_uri": row_dict.get("image_uri", ""),
            # Timestamp
            "created_timestamp": row_dict.get("created_timestamp", ""),
            # Social links (preserved but not used by SimpleLaunchDetectionProcessor)
            "website": row_dict.get("website", ""),
            "twitter": row_dict.get("twitter", ""),
            "telegram": row_dict.get("telegram", ""),
            # Market data (preserved for reference)
            "market_cap": row_dict.get("market_cap", ""),
            "usd_market_cap": row_dict.get("usd_market_cap", ""),
            "total_supply": row_dict.get("total_supply", ""),
            # Include all other CSV columns for completeness
            **row_dict,
        }

        # Create TokenLaunchEvent
        return TokenLaunchEvent(
            raw_data=raw_data,
            timestamp=datetime.now().isoformat(),
            platform="csv_import",
        )

    def _save_failed_row(self, row: Dict[str, str], error: str):
        """Save failed row to a separate file for debugging"""
        failed_file = self.csv_file.with_suffix(".failed.csv")

        try:
            # Add error columns
            row_with_error = row.copy()
            row_with_error["_error"] = error
            row_with_error["_failed_timestamp"] = datetime.now().isoformat()

            # Determine fieldnames for failed CSV
            failed_fieldnames = (
                list(self.fieldnames) if self.fieldnames else list(row.keys())
            )
            if "_error" not in failed_fieldnames:
                failed_fieldnames.extend(["_error", "_failed_timestamp"])

            # Create or append to failed CSV
            file_exists = failed_file.exists()

            with open(
                failed_file, "a" if file_exists else "w", newline="", encoding="utf-8"
            ) as f:
                writer = csv.DictWriter(f, fieldnames=failed_fieldnames)
                if not file_exists:
                    writer.writeheader()
                writer.writerow(row_with_error)

            logger.debug(f"💾 Failed row saved to: {failed_file}")

        except Exception as e:
            logger.error(f"❌ Failed to save failed row: {e}")

    async def start_listening(self):
        """Start listening (compatibility with DefaultWorkflow)"""
        await self.start()

    async def stop(self):
        """Stop the CSV source"""
        self.is_running = False
        logger.info("🛑 Stopping CSV token input source")