import json
from copy import deepcopy
from datetime import datetime, timedelta, UTC
from itertools import groupby
from typing import Dict, Any, List, Tuple, cast

import pytz
from dateutil.parser import parse
from loguru import logger

from bb_integrations_lib.gravitate.pe_api import GravitatePEAPI
from bb_integrations_lib.gravitate.rita_api import GravitateRitaAPI
from bb_integrations_lib.mappers.prices.model import (
    PricePublisher,
)
from bb_integrations_lib.models.pipeline_structs import StopPipeline
from bb_integrations_lib.protocols.pipelines import Step, ParserBase
from bb_integrations_lib.shared.model import PEPriceData


class ImpossibleToParseDate(Exception):
    pass


class PEPriceExportStep(Step):
    def __init__(
        self,
        rita_client: GravitateRitaAPI,
        pe_client: GravitatePEAPI,
        price_publishers: list[PricePublisher],
        config_id: str | None = None,
        parser: type[ParserBase] | None = None,
        parser_kwargs: dict | None = None,
        max_sync_id_override: int | None = None,
        max_sync_datetime_override: datetime | None = None,
        addl_endpoint_args: dict | None = None,
        *args,
        sync_date_tz: pytz.BaseTzInfo,
        **kwargs,
    ):
        """This step requires:
        - price_publishers: [REQUIRED] a list of price publisher names to be included in the price request
        - config_id: [OPTIONAL] RITA config ID to load/store pe_max_sync cursor; defaults to 2 days ago if not stored
        - max_sync_datetime_override: [OPTIONAL] override the MaxSyncDateTime cursor at construction time
        - max_sync_id_override: [OPTIONAL] override the MaxSyncPkId cursor at construction time
        - sync_date_tz: [REQUIRED] timezone the PE expects for MaxSyncDateTime — PE is not timezone-aware and returns naive datetimes in local time, so this is used to localize and convert correctly
        """
        super().__init__(*args, **kwargs)
        self.pe_client = pe_client
        self.price_publishers = price_publishers
        self.rita_client = rita_client
        self.config_id = config_id
        self.additional_endpoint_arguments = addl_endpoint_args or {}
        self.max_sync_datetime_override = max_sync_datetime_override
        self.max_sync_id_override = max_sync_id_override
        self.sync_date_tz = sync_date_tz
        if parser:
            self.custom_parser = parser
            self.custom_parser_kwargs = parser_kwargs or {}

    def price_publisher_lkp(self) -> Dict[str, PricePublisher]:
        lkp = {}
        pp = self.price_publishers
        for p in pp:
            lkp[p.name] = p
        return lkp

    def get_publisher_contract_id(
        self, key: str, actual_contract_id: str | int
    ) -> str | None:
        lkp = self.price_publisher_lkp()
        pub = lkp[key]
        if pub.contract_id_override:
            return pub.contract_id_override
        if pub.use_contract_id:
            return actual_contract_id
        return None

    def get_publisher_extend_by(self, key: str) -> int | None:
        lkp = self.price_publisher_lkp()
        return lkp[key].extend_by_days

    def get_publisher_price_type(self, key: str) -> str:
        lkp = self.price_publisher_lkp()
        return lkp[key].price_type

    def price_type_rows(self, rows: List[PEPriceData]) -> List[PEPriceData]:
        for row in rows:
            price_type = self.get_publisher_price_type(row.PricePublisher)
            row.PriceType = price_type
        return rows

    def describe(self) -> str:
        return "Export Pricing Engine Prices"

    def resolve_max_sync(self) -> Tuple[datetime, int]:
        if self.max_sync_datetime_override is not None:
            return self.max_sync_datetime_override, self.max_sync_id_override or 0

        max_sync = self.pipeline_context.max_sync if self.pipeline_context else None

        if not max_sync:
            raise RuntimeError("No max_sync found in pipeline context")

        pe_max_sync = (max_sync.context or {}).get("pe_max_sync", {})
        max_sync_datetime = pe_max_sync.get("MaxSyncDateTime")

        if max_sync_datetime:
            sync_dt = parse(max_sync_datetime)
        else:
            logger.warning("No pe_max_sync found in context; using default sync date (initial load)")
            sync_dt = datetime.now(UTC) - timedelta(days=1)

        if sync_dt.tzinfo is None:
            sync_dt = cast(pytz.tzinfo.BaseTzInfo, self.sync_date_tz).localize(sync_dt)  # type: ignore[union-attr]
        else:
            sync_dt = sync_dt.astimezone(self.sync_date_tz)

        sync_pk_id = (
            self.max_sync_id_override
            if self.max_sync_id_override is not None
            else pe_max_sync.get("MaxSyncPkId", 0)
        )
        return sync_dt, sync_pk_id

    async def execute(self, _: Any = None) -> List[PEPriceData] | List[Dict]:
        max_sync_datetime, max_sync_pk_id = self.resolve_max_sync()
        updated_prices = await self.get_updated_prices_for_publishers(
            max_sync_datetime=max_sync_datetime,
            max_sync_pk_id=max_sync_pk_id,
            price_publishers=self.price_publishers,
        )
        if not updated_prices:
            logger.info("No updated prices found")
            raise StopPipeline("No updated prices found")

        prices = self.update_historical_prices(updated_prices)
        if not hasattr(self, "custom_parser"):
            return prices
        else:
            parser = self.custom_parser(**self.custom_parser_kwargs)
            parser_results = await parser.parse(prices)
            return parser_results

    def update_historical_prices(self, rows: List[PEPriceData]) -> List[PEPriceData]:
        _sorted_id = sorted(
            rows,
            key=lambda r: (r.PriceInstrumentId, r.EffectiveFromDateTime),
            reverse=True,
        )
        for instrument_id, group in groupby(
            _sorted_id, key=lambda r: r.PriceInstrumentId
        ):
            group_list = PEPriceExportStep.rank_rows(list(group))
            group_list_price_typed = self.price_type_rows(group_list)
            max_row = max(group_list_price_typed, key=lambda r: r.EffectiveFromDateTime)
            max_row.ExtendByDays = self.get_publisher_extend_by(max_row.PricePublisher)
            max_row.ContractId = self.get_publisher_contract_id(
                max_row.PricePublisher, max_row.SourceContractId
            )
            max_row.IsLatest = True
        return _sorted_id

    async def get_prices(
        self, query: Dict, count: int = 1000, include_source_data: bool = True
    ) -> List[PEPriceData]:
        payloads = []
        records = []
        max_sync = None
        payload = {
            "Query": {**query, "COUNT": count},
            "includeSourceData": include_source_data,
        }
        resp = await self.pe_client.get_prices(payload)
        payloads.append(deepcopy(payload))
        json_resp = resp.json()
        while len(json_resp["Data"]) > 0:
            records.extend(json_resp["Data"])
            max_sync = json_resp["MaxSyncResult"]
            if max_sync is None:
                break
            payload["Query"]["MaxSync"] = max_sync
            resp = await self.pe_client.get_prices(payload)
            payloads.append(deepcopy(payload))
            json_resp = resp.json()
        self.pipeline_context.included_files["pe_response"] = json.dumps(records)
        self.pipeline_context.included_files["pe_payloads"] = json.dumps(payloads)
        if max_sync is not None:
            self.pipeline_context.extra_data["pe_max_sync"] = json.dumps(max_sync)
        return [PEPriceData.model_validate(price) for price in records]

    async def get_updated_prices_for_publishers(
        self,
        max_sync_datetime: datetime,
        max_sync_pk_id: int = 0,
        price_publishers: List[PricePublisher] = None,
    ) -> List[PEPriceData]:
        if max_sync_datetime.tzinfo is None:
            max_sync_datetime = max_sync_datetime.replace(tzinfo=pytz.UTC)
        payload = {
            "IsActiveFilterType": "ActiveOnly",
            "PricePublisherNames": [p.name for p in (price_publishers or [])],
            "MaxSync": {
                "MaxSyncDateTime": max_sync_datetime.isoformat(),
                "MaxSyncPkId": max_sync_pk_id,
            },
            **self.additional_endpoint_arguments,
        }
        rows = await self.get_prices(query=payload, include_source_data=True)
        return rows

    @staticmethod
    def rank_rows(rows: List[PEPriceData]) -> List[PEPriceData]:
        for idx, row in enumerate(rows):
            row.Rank = idx + 1
        return rows

    @staticmethod
    def try_to_parse_date(dt_string: str) -> datetime:
        if isinstance(dt_string, str):
            try:
                parsed_datetime = parse(dt_string)
                return parsed_datetime
            except (ValueError, TypeError):
                raise ImpossibleToParseDate(f"Could not parse date: {dt_string}")
        elif isinstance(dt_string, datetime):
            return dt_string
        else:
            raise ImpossibleToParseDate(
                f"Could not parse date: {dt_string} -> Format not supported"
            )
