import concurrent.futures
from jse_tools.core import download_and_process
from .checkpoint import CheckpointManager
from .logger import get_logger


class IngestionEngine:

    def __init__(
        self,
        tickers,
        start_date,
        end_date,
        output_dir="Stocks",
        batch_size=20,
        max_workers=4
    ):

        self.tickers = tickers
        self.start_date = start_date
        self.end_date = end_date
        self.output_dir = output_dir

        self.batch_size = batch_size
        self.max_workers = max_workers

        self.checkpoint = CheckpointManager()
        self.logger = get_logger()

    # ------------------------
    # Batch Generator
    # ------------------------

    def _batch_tickers(self):
        for i in range(0, len(self.tickers), self.batch_size):
            yield self.tickers[i:i + self.batch_size]

    # ------------------------
    # Worker Task
    # ------------------------

    def _process_batch(self, batch):

        try:
            download_and_process(
                batch,
                self.start_date,
                self.end_date,
                self.output_dir
            )

            self.logger.info(f"Batch completed: {batch}")

            return batch

        except Exception as e:
            self.logger.error(f"Batch failure: {e}")
            raise

    # ------------------------
    # Pipeline Runner
    # ------------------------

    def run(self):

        checkpoint_state = self.checkpoint.load_state()

        if checkpoint_state:
            self.logger.info("Resuming from checkpoint")

        batches = list(self._batch_tickers())

        with concurrent.futures.ThreadPoolExecutor(
            max_workers=self.max_workers
        ) as executor:

            executor.map(self._process_batch, batches)

        self.checkpoint.save_state({
            "status": "COMPLETED",
            "total_tickers": len(self.tickers)
        })

        self.logger.info("Pipeline execution finished")
