import io
import logging
import os
import threading
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import date
from threading import current_thread
from typing import (
    List,
    Optional,
)

import pyarrow.parquet as pq
from order_picture_base_model import OrderPicture

from ._Boto3S3StorageService import Boto3S3StorageService
from ._DatalakeConfig import DatalakeConfig
from ._GoldRowToOrderPictureMapper import GoldRowToOrderPictureMapper

logger = logging.getLogger(__name__)


@dataclass
class OutboxRead:
    orders: List[OrderPicture] = field(default_factory=list)
    consumed_paths: List[str] = field(default_factory=list)


class GoldOutboxOrderPictureService:
    def __init__(
        self,
        storage_service: Boto3S3StorageService,
        datalake_config: DatalakeConfig,
    ):
        self.storage_service = storage_service
        self.datalake_config = datalake_config

        cpu_count = os.cpu_count() or 1
        current_context = getattr(threading.current_thread(), "_edp_context_name", None)

        def _context_name_initializer():
            if current_context is not None:
                setattr(threading.current_thread(), "_edp_context_name", current_context)

        self.thread_pool = ThreadPoolExecutor(
            max_workers=20 * cpu_count,
            thread_name_prefix=f"{current_thread().name}$GoldOutboxServiceThread",
            initializer=_context_name_initializer,
        )

    def get_orders_of_store_and_day(
        self,
        store_code: str,
        business_date: date,
    ) -> OutboxRead:
        prefix = self._outbox_day_prefix(store_code, business_date)
        snapshot = self.storage_service.list_files(path=prefix)

        if not snapshot:
            return OutboxRead(orders=[], consumed_paths=[])

        latest_path = max(snapshot)
        logger.info(
            "Reading latest outbox parquet for [%s] [%s]: %s (of %d candidates)",
            store_code,
            business_date,
            latest_path,
            len(snapshot),
        )

        content = self.storage_service.get_file_content(full_path=latest_path)
        rows = pq.ParquetFile(io.BytesIO(content)).read().to_pylist()

        orders = [GoldRowToOrderPictureMapper.map(row) for row in rows]

        return OutboxRead(orders=orders, consumed_paths=snapshot)

    def delete_consumed_files(self, paths: List[str]) -> None:
        if not paths:
            return
        list(self.thread_pool.map(self.storage_service.delete_file, paths))

    def get_first_of_store(self, store_code: str) -> Optional[date]:
        store_prefix = f"{self.datalake_config.outbox_root_path}/storeCode={store_code}/"
        years = self._list_partition_values(path=store_prefix, partition_key="year=")
        if not years:
            return None
        min_year = min(int(y) for y in years)

        year_path = f"{store_prefix}year={min_year:04d}/"
        months = self._list_partition_values(path=year_path, partition_key="month=")
        if not months:
            return None
        min_month = min(int(m) for m in months)

        month_path = f"{year_path}month={min_month:02d}/"
        days = self._list_partition_values(path=month_path, partition_key="day=")
        if not days:
            return None
        min_day = min(int(d) for d in days)

        return date(year=min_year, month=min_month, day=min_day)

    def list_active_store_codes_in_outbox(self) -> List[str]:
        outbox_prefix = f"{self.datalake_config.outbox_root_path}/"
        raw = self.storage_service.list_directories(path=outbox_prefix)
        return [self._strip_partition_prefix(value, "storeCode=") for value in raw]

    def close(self):
        self.thread_pool.shutdown(wait=True, cancel_futures=True)
        self.thread_pool = None

    def _outbox_day_prefix(self, store_code: str, business_date: date) -> str:
        return (
            f"{self.datalake_config.outbox_root_path}/"
            f"storeCode={store_code}/"
            f"year={business_date.year:04d}/"
            f"month={business_date.month:02d}/"
            f"day={business_date.day:02d}/"
        )

    def _list_partition_values(self, path: str, partition_key: str) -> List[str]:
        raw = self.storage_service.list_directories(path=path)
        return [self._strip_partition_prefix(value, partition_key) for value in raw]

    @staticmethod
    def _strip_partition_prefix(value: str, prefix: str) -> str:
        if value.startswith(prefix):
            return value[len(prefix):]
        return value
