import os
import time
import logging
import json
import boto3

from typing import Any, Dict, List, Optional, Tuple
from botocore.exceptions import ClientError
from boto3.dynamodb.conditions import Key, Attr

logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

AWS_REGION = os.getenv("AWS_DEFAULT_REGION", "us-east-1")

class BaseDynamoDBRepository:
    """
    Python equivalent of your TypeScript BaseDynamoDBRepository.
    Use by subclassing and setting:
      self.TABLE_NAME = "your-table"
      self.PRIMARY_KEY = ["pk", "sk"]  # optional keys used for KeyCondition
    """

    def __init__(self, region: str = AWS_REGION):
        self.TABLE_NAME: str = ""
        self.PRIMARY_KEY: List[str] = []
        self._dynamodb = boto3.resource("dynamodb", region_name=region)

    @property
    def table(self):
        if not self.TABLE_NAME:
            raise RuntimeError("TABLE_NAME not set on repository")
        return self._dynamodb.Table(self.TABLE_NAME)

    # -----------------------
    # Create / Update / Delete
    # -----------------------
    def add(self, item: Dict[str, Any]) -> Dict[str, Any]:
        """Put item into table. Returns the item (similar to your TS)"""
        try:
            item = self.parse_date(item)
            now = int(time.time() * 1000)
            item.setdefault("createdAt", now)
            item.setdefault("updatedAt", now)
            self.table.put_item(Item=item)
            logger.info("dbItem added")
            return item
        except ClientError as err:
            self.handle_error(err)
            return {}

    def update(self, key_name: str, item: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Update a single-key item. key_name is the attribute used as the partition key in `item`.
        Example: key_name="id", item={"id": "123", "count": 10, "name": "x"}
        """
        try:
            # Build attribute names/values and update expression
            item_keys = [k for k in item.keys() if k != key_name]
            if not item_keys:
                return None

            expr_names = {f"#field{i}": k for i, k in enumerate(item_keys)}
            expr_values = {f":value{i}": item[k] for i, k in enumerate(item_keys)}
            update_expr = "SET " + ", ".join([f"#field{i} = :value{i}" for i in range(len(item_keys))])

            key = {key_name: item[key_name]}
            resp = self.table.update_item(
                Key=key,
                UpdateExpression=update_expr,
                ExpressionAttributeNames=expr_names,
                ExpressionAttributeValues=expr_values,
                ReturnValues="ALL_NEW",
            )
            return resp.get("Attributes")
        except ClientError as err:
            self.handle_error(err)
            return None

    def update_by_multiple(self, keys: List[str], item: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Update using multiple key attributes. `keys` is list of attribute names that form the Key.
        `item` must contain values for those key attributes.
        """
        try:
            item_keys = [k for k in item.keys() if k not in keys]
            if not item_keys:
                return None

            expr_names = {f"#field{i}": k for i, k in enumerate(item_keys)}
            expr_values = {f":value{i}": item[k] for i, k in enumerate(item_keys)}
            update_expr = "SET " + ", ".join([f"#field{i} = :value{i}" for i in range(len(item_keys))])

            key_obj = {k: item[k] for k in keys}
            resp = self.table.update_item(
                Key=key_obj,
                UpdateExpression=update_expr,
                ExpressionAttributeNames=expr_names,
                ExpressionAttributeValues=expr_values,
                ReturnValues="ALL_NEW",
            )
            return resp.get("Attributes")
        except ClientError as err:
            self.handle_error(err)
            return None

    def update_by_multiple_with_list_append(self, keys: List[str], item: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        Update where some fields are lists and should be appended using list_append.
        """
        try:
            item_keys = [k for k in item.keys() if k not in keys]
            if not item_keys:
                return None

            expr_names = {}
            expr_values = {}
            update_parts = []

            for i, k in enumerate(item_keys):
                expr_names[f"#field{i}"] = k
                placeholder = f":value{i}"
                expr_values[placeholder] = item[k]

                # if list -> use list_append
                if isinstance(item[k], list):
                    update_parts.append(f"#field{i} = list_append(if_not_exists(#field{i}, :empty_list), {placeholder})")
                    # ensure empty_list exists
                    expr_values[":empty_list"] = expr_values.get(":empty_list", [])
                else:
                    update_parts.append(f"#field{i} = {placeholder}")

            update_expr = "SET " + ", ".join(update_parts)

            key_obj = {k: item[k] for k in keys}
            resp = self.table.update_item(
                Key=key_obj,
                UpdateExpression=update_expr,
                ExpressionAttributeNames=expr_names,
                ExpressionAttributeValues=expr_values,
                ReturnValues="ALL_NEW",
            )
            return resp.get("Attributes")
        except ClientError as err:
            self.handle_error(err)
            return None

    # -----------------------
    # Read operations
    # -----------------------
    def get_by_id(self, key: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """Get single item by key object"""
        try:
            resp = self.table.get_item(Key=key)
            item = resp.get("Item")
            if item:
                return self.unparse_date(item)
            return None
        except ClientError as err:
            logger.error("error occured while getting item by primary keys - %s", err)
            self.handle_error(err)
            return None

    # The query implementation below supports basic comparators and IN/BETWEEN.
    def query(
        self,
        item: Dict[str, Any],
        condition: Optional[List[str]] = None,
        index_name: Optional[str] = None,
        projection_expression: Optional[str] = None,
        limit: Optional[int] = None,
        scan_index_forward: bool = True,
        exclusive_start_key: Optional[Dict[str, Any]] = None,
    ) -> Optional[List[Dict[str, Any]]]:
        """
        Performs a query. `item` is mapping of attribute values to query.
        `condition` is optional list of comparator strings for each attribute (e.g. ['=', 'IN', 'BETWEEN']).
        This function builds KeyConditionExpression for primary keys (or index keys if index_name provided),
        and uses Attr filter for non-key attributes.
        """
        try:
            item_keys = list(item.keys())
            # create KeyCondition for primary key attributes (or index_name split if provided)
            key_cond = None
            filter_cond = None

            # helper to get comparator for index
            def comp_at(idx):
                return (condition[idx] if condition and idx < len(condition) else "=").upper()

            # Build KeyConditionExpression using PRIMARY_KEY or index_name if provided
            key_names = self.PRIMARY_KEY if self.PRIMARY_KEY else [item_keys[0]]
            key_expr = None
            for idx, k in enumerate(item_keys):
                if (index_name and (k in index_name.split("-"))) or (not index_name and k in key_names):
                    comp = comp_at(idx)
                    val = item[k]
                    if comp == "=":
                        expr = Key(k).eq(val)
                    elif comp == "BETWEEN" and isinstance(val, (list, tuple)) and len(val) == 2:
                        expr = Key(k).between(val[0], val[1])
                    elif comp == ">" :
                        expr = Key(k).gt(val)
                    elif comp == ">=" :
                        expr = Key(k).gte(val)
                    elif comp == "<" :
                        expr = Key(k).lt(val)
                    elif comp == "<=" :
                        expr = Key(k).lte(val)
                    elif comp == "BEGINS_WITH":
                        # boto3 Key has begins_with via Key(k).begins_with
                        expr = Key(k).begins_with(val)
                    elif comp == "IN" and isinstance(val, (list, tuple)):
                        # Key does not support IN directly; we fallback to multiple 'or' keys via filter later.
                        expr = Key(k).eq(val[0])
                    else:
                        expr = Key(k).eq(val)

                    key_expr = expr if key_expr is None else key_expr & expr

            # For non-key attributes build FilterExpression using Attr
            for idx, k in enumerate(item_keys):
                if (index_name and (k in index_name.split("-"))) or (not index_name and k in key_names):
                    continue
                comp = comp_at(idx)
                val = item[k]
                if comp == "=":
                    fexpr = Attr(k).eq(val)
                elif comp == "IN" and isinstance(val, (list, tuple)):
                    # build OR of equality operations
                    fexpr = None
                    for v in val:
                        term = Attr(k).eq(v)
                        fexpr = term if fexpr is None else fexpr | term
                elif comp == "BETWEEN" and isinstance(val, (list, tuple)) and len(val) == 2:
                    fexpr = Attr(k).between(val[0], val[1])
                elif comp == "contains":
                    fexpr = Attr(k).contains(val)
                else:
                    # fallback to eq
                    fexpr = Attr(k).eq(val)

                filter_cond = fexpr if filter_cond is None else filter_cond & fexpr

            # build query params
            params: Dict[str, Any] = {"ScanIndexForward": scan_index_forward}
            if key_expr is not None:
                params["KeyConditionExpression"] = key_expr
            if filter_cond is not None:
                params["FilterExpression"] = filter_cond
            if projection_expression:
                params["ProjectionExpression"] = projection_expression
            if exclusive_start_key:
                params["ExclusiveStartKey"] = exclusive_start_key
            if index_name:
                params["IndexName"] = index_name

            acc: List[Dict[str, Any]] = []
            resp = self.table.query(**params)
            acc.extend(resp.get("Items", []))
            while "LastEvaluatedKey" in resp and (not limit or len(acc) < limit):
                params["ExclusiveStartKey"] = resp["LastEvaluatedKey"]
                if limit:
                    # adjust client-side limiting by leaving unchanged but truncating results below
                    pass
                resp = self.table.query(**params)
                acc.extend(resp.get("Items", []))

            if not acc:
                return None
            if limit:
                acc = acc[:limit]
            return self.unparse_date(acc)
        except ClientError as err:
            self.handle_error(err)
            return None

    def query_with_pagination(
        self,
        item: Dict[str, Any],
        condition: Optional[List[str]] = None,
        index_name: Optional[str] = None,
        projection_expression: Optional[str] = None,
        limit: int = 1000,
        scan_index_forward: bool = True,
        exclusive_start_key: Optional[Dict[str, Any]] = None,
    ) -> Optional[Dict[str, Any]]:
        """
        Similar to query, but returns { 'exclusiveStartKey': key, 'data': [...] } and enforces limit
        """
        try:
            # We'll reuse query but manually control pagination
            params: Dict[str, Any] = {}
            # Reuse query logic but do manual loop to fill up to limit
            item_keys = list(item.keys())

            # Build expressions as in query()
            key_names = self.PRIMARY_KEY if self.PRIMARY_KEY else [item_keys[0]]
            key_expr = None
            filter_cond = None

            def comp_at(idx):
                return (condition[idx] if condition and idx < len(condition) else "=").upper()

            for idx, k in enumerate(item_keys):
                if (index_name and (k in index_name.split("-"))) or (not index_name and k in key_names):
                    comp = comp_at(idx)
                    val = item[k]
                    if comp == "=":
                        expr = Key(k).eq(val)
                    elif comp == "BETWEEN" and isinstance(val, (list, tuple)) and len(val) == 2:
                        expr = Key(k).between(val[0], val[1])
                    else:
                        expr = Key(k).eq(val)
                    key_expr = expr if key_expr is None else key_expr & expr

            for idx, k in enumerate(item_keys):
                if (index_name and (k in index_name.split("-"))) or (not index_name and k in key_names):
                    continue
                comp = comp_at(idx)
                val = item[k]
                if comp == "=":
                    fexpr = Attr(k).eq(val)
                elif comp == "IN" and isinstance(val, (list, tuple)):
                    fexpr = None
                    for v in val:
                        term = Attr(k).eq(v)
                        fexpr = term if fexpr is None else fexpr | term
                elif comp == "BETWEEN" and isinstance(val, (list, tuple)) and len(val) == 2:
                    fexpr = Attr(k).between(val[0], val[1])
                else:
                    fexpr = Attr(k).eq(val)
                filter_cond = fexpr if filter_cond is None else filter_cond & fexpr

            params = {"KeyConditionExpression": key_expr, "ScanIndexForward": scan_index_forward}
            if filter_cond:
                params["FilterExpression"] = filter_cond
            if index_name:
                params["IndexName"] = index_name
            if projection_expression:
                params["ProjectionExpression"] = projection_expression
            if exclusive_start_key:
                params["ExclusiveStartKey"] = exclusive_start_key
            if limit:
                params["Limit"] = limit

            acc: Dict[str, Any] = {"exclusiveStartKey": None, "data": []}
            resp = self.table.query(**params)
            acc["data"].extend(resp.get("Items", []))
            acc["exclusiveStartKey"] = resp.get("LastEvaluatedKey")
            while resp.get("LastEvaluatedKey") and len(acc["data"]) < limit:
                params["ExclusiveStartKey"] = resp["LastEvaluatedKey"]
                # adjust limit for this iteration
                resp = self.table.query(**params)
                acc["data"].extend(resp.get("Items", []))
                acc["exclusiveStartKey"] = resp.get("LastEvaluatedKey")

            if not acc["data"]:
                return None
            acc["data"] = self.unparse_date(acc["data"])
            return acc
        except ClientError as err:
            self.handle_error(err)
            return None

    def scan(
        self,
        item: Dict[str, Any],
        condition: List[str],
        projection_expression: Optional[str] = None,
        exclusive_start_key: Optional[Dict[str, Any]] = None,
    ) -> Optional[List[Dict[str, Any]]]:
        """
        Scan with filter expressions. `condition` must be same length as keys in item.
        """
        try:
            item_keys = list(item.keys())
            # Build filter expression
            filter_cond = None
            for idx, k in enumerate(item_keys):
                comp = (condition[idx] if condition and idx < len(condition) else "=").upper()
                val = item[k]
                if comp == "=":
                    fexpr = Attr(k).eq(val)
                elif comp == "IN" and isinstance(val, (list, tuple)):
                    fexpr = None
                    for v in val:
                        term = Attr(k).eq(v)
                        fexpr = term if fexpr is None else fexpr | term
                elif comp == "BETWEEN" and isinstance(val, (list, tuple)) and len(val) == 2:
                    fexpr = Attr(k).between(val[0], val[1])
                else:
                    fexpr = Attr(k).eq(val)

                filter_cond = fexpr if filter_cond is None else filter_cond & fexpr

            params: Dict[str, Any] = {}
            if filter_cond:
                params["FilterExpression"] = filter_cond
            if projection_expression:
                params["ProjectionExpression"] = projection_expression
            if exclusive_start_key:
                params["ExclusiveStartKey"] = exclusive_start_key

            acc: List[Dict[str, Any]] = []
            resp = self.table.scan(**params)
            acc.extend(resp.get("Items", []))
            while "LastEvaluatedKey" in resp:
                params["ExclusiveStartKey"] = resp["LastEvaluatedKey"]
                resp = self.table.scan(**params)
                acc.extend(resp.get("Items", []))

            if not acc:
                return None
            return self.unparse_date(acc)
        except ClientError as err:
            self.handle_error(err)
            return None

    def delete(self, key: Dict[str, Any]) -> None:
        try:
            self.table.delete_item(Key=key)
            logger.info("Dynamo db Repository: delete End")
        except ClientError as err:
            logger.error("error occured while deleting item - %s", err)
            self.handle_error(err)

    def get_all(self, exclusive_start_key: Optional[Dict[str, Any]] = None) -> List[Dict[str, Any]]:
        params: Dict[str, Any] = {}
        if exclusive_start_key:
            params["ExclusiveStartKey"] = exclusive_start_key
        acc: List[Dict[str, Any]] = []
        try:
            resp = self.table.scan(**params)
            acc.extend(resp.get("Items", []))
            while "LastEvaluatedKey" in resp:
                params["ExclusiveStartKey"] = resp["LastEvaluatedKey"]
                resp = self.table.scan(**params)
                acc.extend(resp.get("Items", []))
            return self.unparse_date(acc)
        except ClientError as err:
            self.handle_error(err)
            return []

    def bulk_upload(self, items: List[Dict[str, Any]]) -> bool:
        """
        Bulk upload using Table.batch_writer() which handles chunking and retries automatically.
        """
        try:
            # parse dates / prepare
            insert_items = [self.parse_date(i) for i in items]
            with self.table.batch_writer(overwrite_by_pkeys=[]) as batch:
                # boto3 batch_writer doesn't accept overwrite_by_pkeys prior to certain versions,
                # leaving default behaviour (it will batch up to 25 per request).
                for it in insert_items:
                    batch.put_item(Item=it)
            logger.info("Upload Successfully")
            return True
        except ClientError as err:
            self.handle_error(err)
            return False

    # -----------------------
    # Utility helpers
    # -----------------------
    def parse_date(self, data: Any) -> Any:
        """Return JSON-serializable copy (mirrors TS parseDate which did JSON.parse/JSON.stringify)"""
        return json.loads(json.dumps(data, default=str))

    def unparse_date(self, data: Any) -> Any:
        """
        Attempt to convert ISO-like date strings back to datetime where applicable.
        For simplicity we only try `.fromisoformat` and fallback if fails.
        """
        # If a list of items
        def reviver(value):
            if isinstance(value, str):
                try:
                    # python fromisoformat supports many ISO variants (without Z).
                    # try replacing trailing Z with +00:00
                    txt = value
                    if txt.endswith("Z"):
                        txt = txt[:-1] + "+00:00"
                    return __import__("datetime").datetime.fromisoformat(txt)
                except Exception:
                    return value
            elif isinstance(value, dict):
                return {k: reviver(v) for k, v in value.items()}
            elif isinstance(value, list):
                return [reviver(v) for v in value]
            else:
                return value

        if isinstance(data, list):
            return [reviver(item) for item in data]
        elif isinstance(data, dict):
            return reviver(data)
        else:
            return data

    def handle_error(self, err: Any) -> None:
        """Map some common DynamoDB errors to helpful messages"""
        if not err:
            logger.error("Encountered error object was empty")
            return
        # botocore ClientError contains response with Error.Code and Error.Message
        code = None
        message = None
        try:
            code = err.response["Error"].get("Code")
            message = err.response["Error"].get("Message")
        except Exception:
            # fallback
            code = getattr(err, "code", None)
            message = str(err)

        if not code:
            logger.error("An exception occurred, investigate and configure retry strategy. Error: %s", message)
            return

        if code == "ConditionalCheckFailedException":
            logger.error("Condition check failed. Error: %s", message)
        elif code == "TransactionConflictException":
            logger.error("Transaction conflict - safe to retry with backoff. Error: %s", message)
        elif code == "ItemCollectionSizeLimitExceededException":
            logger.error("Item collection too large - consider GSI. Error: %s", message)
        else:
            # common errors
            self.handle_common_errors(code, message)

    def handle_common_errors(self, code: str, message: str) -> None:
        if code == "InternalServerError":
            logger.error("Internal Server Error. Error: %s", message)
        elif code == "ProvisionedThroughputExceededException":
            logger.error("Throughput exceeded - consider retries or higher capacity. Error: %s", message)
        elif code == "ResourceNotFoundException":
            logger.error("Table not found. Error: %s", message)
        elif code == "ServiceUnavailable":
            logger.error("Service unavailable. Error: %s", message)
        elif code == "ThrottlingException":
            logger.error("Request throttled. Error: %s", message)
        elif code == "UnrecognizedClientException":
            logger.error("Invalid credentials. Error: %s", message)
        elif code == "ValidationException":
            logger.error("Validation error. Error: %s", message)
        elif code == "RequestLimitExceeded":
            logger.error("Account request limit exceeded. Error: %s", message)
        else:
            logger.error("Unhandled DynamoDB error: %s - %s", code, message)

    @staticmethod
    def is_comparator(s: Any) -> bool:
        return s in ["=", "<>", "<", "<=", ">", ">="]

    @staticmethod
    def get_array_value(item: List[Any], index: int) -> Dict[str, Any]:
        """Return expression attribute values mapping for an array (similar idea used in TS)"""
        return {f":value{index}{idx}": v for idx, v in enumerate(item)}

    def get_key_and_filter(self, item: Dict[str, Any], condition: Optional[List[str]] = None, index_name: Optional[str] = None) -> Tuple[str, str]:
        """
        Helper to produce Key and Filter expression strings similar to your TS method.
        In Python we typically use Key()/Attr() objects; this function kept for compatibility.
        Returns tuple (KeyExpressionString, FilterExpressionString).
        NOTE: In this Python port the main query() and scan() use condition-building through Key/Attr objects,
              so this helper is provided for potential low-level use and testing.
        """
        try:
            item_keys = list(item.keys())
            # Build simplified string forms (not used by default query implementation)
            key_parts = []
            filter_parts = []

            def comp_at(idx):
                return (condition[idx] if condition and idx < len(condition) else "=").upper()

            for idx, k in enumerate(item_keys):
                comp = comp_at(idx)
                if not index_name:
                    if self.PRIMARY_KEY and k in self.PRIMARY_KEY or (not self.PRIMARY_KEY):
                        if comp == "IN" and isinstance(item[k], (list, tuple)):
                            vals = ",".join([str(x) for x in item[k]])
                            key_parts.append(f"{k} IN ({vals})")
                        elif comp == "BETWEEN" and isinstance(item[k], (list, tuple)) and len(item[k]) == 2:
                            key_parts.append(f"{k} BETWEEN {item[k][0]} AND {item[k][1]}")
                        else:
                            key_parts.append(f"{k} {comp} :value{idx}")
                    else:
                        if comp == "IN" and isinstance(item[k], (list, tuple)):
                            vals = ",".join([str(x) for x in item[k]])
                            filter_parts.append(f"{k} IN ({vals})")
                        elif comp == "BETWEEN" and isinstance(item[k], (list, tuple)) and len(item[k]) == 2:
                            filter_parts.append(f"{k} BETWEEN {item[k][0]} AND {item[k][1]}")
                        else:
                            filter_parts.append(f"{k} {comp} :value{idx}")
                else:
                    # if index_name provided, split and use that as primary fields
                    index_fields = index_name.split("-")
                    if k in index_fields:
                        key_parts.append(f"{k} {comp} :value{idx}")
                    else:
                        filter_parts.append(f"{k} {comp} :value{idx}")

            KeyExpr = " and ".join([p for p in key_parts if p])
            FilterExpr = " and ".join([p for p in filter_parts if p])
            return KeyExpr, FilterExpr
        except Exception as e:
            raise

