"""
Core DigitCounter implementation – fully rewritten and corrected.
"""

import math
import time
from typing import Any, Callable, Dict, List, Optional, Union
from functools import lru_cache
from decimal import Decimal, getcontext

from .constants import (
    DEFAULT_PRECISION,
    MAX_PRECISION,
    CACHE_SIZE,
    SCIENTIFIC_NOTATION_CHARS,
    DECIMAL_POINT,
)
from .exceptions import (
    DigitCounterError,
    InputValidationError,
    CalculationError,
    CacheError,
    BatchProcessingError,
)
from .types import DigitCountResult, BatchResult, CountingMode, LogLevel
from .utils import (
    validate_precision,
    format_number,
    parse_scientific_notation,
    setup_logging,
)


# =============================================================
# DigitCounter
# =============================================================

class DigitCounter:
    """High-performance digit counter (int, float, Decimal, scientific)."""

    def __init__(
        self,
        precision: int = DEFAULT_PRECISION,
        enable_cache: bool = True,
        log_level: Union[int, str, LogLevel] = LogLevel.INFO,
        cache_size: int = CACHE_SIZE,
        warn_on_precision_limit: bool = True,
        use_decimal_for_all: bool = True,
    ):
        """Initialize counter."""
        if isinstance(log_level, LogLevel):
            log_level = log_level.value

        self.logger = setup_logging(level=log_level)
        self.precision = validate_precision(precision)

        if warn_on_precision_limit and self.precision < precision:
            self.logger.warning(
                "Precision adjusted from %d to %d (max=%d)",
                precision, self.precision, MAX_PRECISION
            )

        self.enable_cache = enable_cache
        self.cache_size = int(cache_size)
        self.use_decimal_for_all = bool(use_decimal_for_all)

        # Decimal precision (extra buffer)
        getcontext().prec = self.precision + 10

        # Cache init
        if self.enable_cache:
            self._init_cache()

        # Statistics
        self._stats = {
            "total_calls": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "total_processing_time_ms": 0.0,
        }

        self.logger.info(
            "DigitCounter initialized | precision=%d | cache=%s | cache_size=%d",
            self.precision,
            "enabled" if enable_cache else "disabled",
            self.cache_size,
        )

    # =============================================================
    # Cache setup
    # =============================================================

    def _init_cache(self) -> None:
        try:
            self._cached_count_digits = lru_cache(maxsize=self.cache_size)(
                self._cached_impl
            )
        except Exception as exc:
            raise CacheError(f"Cache initialization failed: {exc}") from exc

    # =============================================================
    # Input validation
    # =============================================================

    def _validate_and_normalize_input(self, value: Any) -> Union[int, float, Decimal]:
        """Validate input and convert to canonical internal type."""

        if isinstance(value, (int, float, Decimal)):
            return value

        if isinstance(value, str):
            raw = value.strip()
            if not raw:
                raise InputValidationError("Empty string", value=value)

            low = raw.lower()
            if low in ("inf", "+inf", "-inf", "infinity", "-infinity"):
                raise InputValidationError("Infinity is not supported", value=value)
            if low in ("nan", "+nan", "-nan"):
                raise InputValidationError("NaN is not supported", value=value)

            # Prefer Decimal (preserves exact form)
            try:
                return Decimal(raw)
            except Exception:
                try:
                    return float(raw)
                except Exception as exc:
                    raise InputValidationError(
                        "Cannot parse string", value=value, original_exception=exc
                    )

        raise InputValidationError(
            f"Unsupported input type: {type(value)}", value=value
        )

    # =============================================================
    # Fast integer digit counting
    # =============================================================

    @staticmethod
    def _count_integer_digits_fast(n: int) -> int:
        """Fast digit count for positive integers."""
        if n < 0:
            n = -n
        if n == 0:
            return 1

        thresholds = [
            10, 100, 1000, 10000, 100000, 1000000,
            10000000, 100000000, 1000000000,
        ]
        for i, t in enumerate(thresholds, start=1):
            if n < t:
                return i
        return int(math.log10(n)) + 1

    # =============================================================
    # Decimal precise counting (NO normalization)
    # =============================================================

    def _count_digits_decimal(self, dec: Decimal) -> Dict[str, int]:
        """Count digits for Decimal without removing trailing zeros."""
        try:
            s = format_number(dec)
            # print(f'float number check {s}')
            # Strip sign
            negative = s.startswith("-")
            if negative:
                s = s[1:]

            # Scientific notation?
            if any(ch in s for ch in SCIENTIFIC_NOTATION_CHARS):
                return self._count_scientific_notation_digits(s)

            # Float-like decimal
            if DECIMAL_POINT in s:
                int_str, frac_str = s.split(DECIMAL_POINT, 1)
                int_str = int_str or "0"
                
                before = self._count_integer_digits_fast(abs(int(int_str)))
                after = len(frac_str)  # EXACT count — trailing zeros included

            else:
                before = self._count_integer_digits_fast(abs(int(s)))
                after = 0

            return {
                "total": before + after,
                "before": before,
                "after": after,
                "is_integer": after == 0,
            }

        except Exception as exc:
            raise CalculationError(
                "Decimal digit counting failed",
                value=dec,
                original_exception=exc,
            )

    # =============================================================
    # Scientific notation handling
    # =============================================================

    def _count_scientific_notation_digits(self, s: str) -> Dict[str, int]:
        try:
            mantissa, exponent = parse_scientific_notation(s)
            expanded = mantissa * (Decimal(10) ** exponent)
            return self._count_digits_decimal(expanded)
        except Exception as exc:
            raise CalculationError(
                "Failed to parse scientific notation",
                value=s,
                original_exception=exc,
            )

    # =============================================================
    # Cache wrapper
    # =============================================================

    def _cached_impl(self, value_str: str, precision: int, mode: str) -> Dict[str, Any]:
        try:
            if any(ch in value_str for ch in SCIENTIFIC_NOTATION_CHARS) or DECIMAL_POINT in value_str:
                value_obj = Decimal(value_str)
            else:
                value_obj = int(value_str)
        except Exception as exc:
            raise CacheError("Cache parse error", value=value_str, original_exception=exc)

        data = self._count_digits_uncached(value_obj, precision)
        data.update({"value_str": value_str, "precision": precision, "mode": mode})
        return data

    # =============================================================
    # Core dispatch
    # =============================================================

    def _count_digits_uncached(self, value: Union[int, float, Decimal], precision: int) -> Dict[str, Any]:
        if isinstance(value, int) or (isinstance(value, float) and value.is_integer()):
            return self._count_integer(value)

        if isinstance(value, Decimal) or self.use_decimal_for_all:
            return self._count_digits_decimal(Decimal(str(value)))

        return self._count_float_string(float(value), precision)

    # =============================================================
    # Integer logic
    # =============================================================

    def _count_integer(self, value: Union[int, float]) -> Dict[str, Any]:
        """Integer rule: after-decimal = before-decimal."""
        n = abs(int(value))
        
        before = self._count_integer_digits_fast(n)
        after = before
        return {"total": before + after, "before": before, "after": after, "is_integer": True}

    # =============================================================
    # Float string mode
    # =============================================================

    def _count_float_string(self, value: float, precision: int) -> Dict[str, Any]:
        s = format_number(value)
        
        if any(ch in s for ch in SCIENTIFIC_NOTATION_CHARS):
            return self._count_scientific_notation_digits(s)

        parts = s.split(DECIMAL_POINT, 1)

        # No decimal point → treat like integer
        if len(parts) == 1:
            before = self._count_integer_digits_fast(abs(int(value)))
            after = before
            return {
                "total": before + after,
                "before": before,
                "after": after,
                "is_integer": True,
            }

        int_str, frac_str = parts
        before = self._count_integer_digits_fast(abs(int(float(int_str)))) if int_str else 0

        # EXACT COUNT — keep trailing zeros
        after = min(len(frac_str), precision)

        return {
            "total": before + after,
            "before": before,
            "after": after,
            "is_integer": after == 0,
        }

    # =============================================================
    # Public count_digits()
    # =============================================================

    def count_digits(
        self,
        value: Union[int, float, str, Decimal],
        mode: Union[str, CountingMode] = CountingMode.TOTAL,
        precision: Optional[int] = None,
        return_dict: bool = False,
    ):
        start = time.perf_counter()
        self._stats["total_calls"] += 1

        mode_str = mode.value if isinstance(mode, CountingMode) else str(mode)

        if mode_str not in CountingMode.values():
            raise InputValidationError(f"Invalid mode: {mode_str}")

        prec = validate_precision(precision) if precision else self.precision
        validated = self._validate_and_normalize_input(value)

        if isinstance(validated, Decimal):
            key = format_number(validated)
        elif isinstance(validated, float):
            key = repr(validated)
        else:
            key = str(validated)

        if self.enable_cache:
            try:
                before = self._cached_count_digits.cache_info()
                data = self._cached_count_digits(key, prec, mode_str)
                after = self._cached_count_digits.cache_info()

                self._stats["cache_hits"] += after.hits - before.hits
                self._stats["cache_misses"] += after.misses - before.misses

            except Exception:
                data = self._count_digits_uncached(validated, prec)
                self._stats["cache_misses"] += 1
        else:
            data = self._count_digits_uncached(validated, prec)

        elapsed = (time.perf_counter() - start) * 1000.0
        self._stats["total_processing_time_ms"] += elapsed

        if mode_str == CountingMode.TOTAL.value:
            return data["total"]
        if mode_str == CountingMode.BEFORE.value:
            return data["before"]
        if mode_str == CountingMode.AFTER.value:
            return data["after"]

        result = DigitCountResult(
            value=value,
            total_digits=data["total"],
            before_decimal=data["before"],
            after_decimal=data["after"],
            precision_used=prec,
            is_integer=data.get("is_integer", False),
            processing_time_ms=elapsed,
        )
        return result.to_dict() if return_dict else result

    # =============================================================
    # Batch processing
    # =============================================================

    def batch_process(
        self,
        values: List[Any],
        mode: Union[str, CountingMode] = CountingMode.TOTAL,
        precision: Optional[int] = None,
        return_dicts: bool = False,
        skip_errors: bool = True,
        progress_callback: Optional[Callable[[int, int], None]] = None,
    ) -> BatchResult:

        start = time.perf_counter()
        mode_str = mode.value if isinstance(mode, CountingMode) else mode
        precision_used = validate_precision(precision) if precision else self.precision

        results = []
        errors = []

        for idx, val in enumerate(values):
            try:
                res = self.count_digits(val, mode_str, precision_used, return_dicts)
                results.append(res)

                if progress_callback:
                    progress_callback(idx + 1, len(values))

            except DigitCounterError as exc:
                err = {"index": idx, "value": val, "error": str(exc)}
                errors.append(err)

                if skip_errors:
                    results.append(None)
                else:
                    raise BatchProcessingError(
                        f"Batch failed at index {idx}",
                        value=val,
                        original_exception=exc,
                    )

        total_ms = (time.perf_counter() - start) * 1000.0
        success_count = sum(1 for r in results if r is not None)

        return BatchResult(
            total_items=len(values),
            successful_items=success_count,
            failed_items=len(errors),
            total_processing_time_ms=total_ms,
            average_time_per_item_ms=(total_ms / len(values)) if values else 0,
            results=results,
            errors=errors,
        )

    # =============================================================
    # Cache & stats
    # =============================================================

    def clear_cache(self) -> None:
        if self.enable_cache:
            self._cached_count_digits.cache_clear()
            self.logger.info("Cache cleared")

    def get_statistics(self) -> Dict[str, Any]:
        stats = self._stats.copy()

        if self.enable_cache:
            info = self._cached_count_digits.cache_info()
            stats.update(
                {
                    "cache_size": self.cache_size,
                    "cache_hits": info.hits,
                    "cache_misses": info.misses,
                    "cache_currsize": info.currsize,
                    "cache_maxsize": info.maxsize,
                    "cache_hit_rate": (
                        info.hits / (info.hits + info.misses) * 100
                        if (info.hits + info.misses)
                        else 0.0
                    ),
                }
            )

        if stats["total_calls"]:
            stats["avg_processing_time_ms"] = (
                stats["total_processing_time_ms"] / stats["total_calls"]
            )

        return stats

    def reset_statistics(self) -> None:
        self._stats = {
            "total_calls": 0,
            "cache_hits": 0,
            "cache_misses": 0,
            "total_processing_time_ms": 0.0,
        }
        self.logger.info("Statistics reset")

    # =============================================================
    # Context manager
    # =============================================================

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        if exc_type:
            self.logger.error("Context manager error: %s", exc)
        return False
