from __future__ import annotations

import itertools
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Iterator

from mathstream import StreamNumber, add, sub, mul, div
from mathstream.number import LOG_DIR


@dataclass
class PiApproximation:
    iteration: int
    digits: str


def _literal(value: int | str) -> StreamNumber:
    return StreamNumber(literal=str(value))


_SCALE_CACHE: dict[int, Path] = {}


def _scaled_one(scale_power: int) -> StreamNumber:
    if scale_power < 0:
        raise ValueError("scale_power must be non-negative")
    cached = _SCALE_CACHE.get(scale_power)
    if cached and cached.exists():
        return StreamNumber(cached)

    LOG_DIR.mkdir(parents=True, exist_ok=True)
    path = LOG_DIR / f"scale_1e{scale_power}.txt"
    if not path.exists():
        with open(path, "w", encoding="utf-8") as handle:
            handle.write("1")
            remaining = scale_power
            if remaining > 0:
                chunk_size = 1_000_000
                zero_chunk = "0" * chunk_size
                while remaining > 0:
                    write_size = min(chunk_size, remaining)
                    handle.write(zero_chunk[:write_size])
                    remaining -= write_size
    _SCALE_CACHE[scale_power] = path
    return StreamNumber(path)


def _format_scaled(value: StreamNumber, scale_power: int, digits: int) -> str:
    raw = "".join(value.stream())
    negative = raw.startswith("-")
    if negative:
        raw = raw[1:]
    if len(raw) <= scale_power:
        raw = raw.zfill(scale_power + 1)
    integer_part = raw[:-scale_power] or "0"
    fractional_part = raw[-scale_power:] if scale_power else ""
    if digits >= 0 and fractional_part:
        fractional_part = fractional_part[:digits]
    text = integer_part
    if digits != 0:
        text = f"{integer_part}.{fractional_part or '0'}"
    if negative:
        text = f"-{text}"
    return text


def iterate_nilakantha(
    iterations: int | None,
    *,
    digits: int,
    extra_precision: int = 4,
) -> Iterator[PiApproximation]:
    """Yield successive Nilakantha approximations of π using mathstream primitives."""
    if iterations is not None and iterations < 0:
        raise ValueError("iterations must be non-negative")
    if digits < 0:
        raise ValueError("digits must be non-negative")

    scale_power = digits + extra_precision
    scale = _scaled_one(scale_power)
    three = _literal(3)
    four = _literal(4)

    total = mul(three, scale)
    numerator = mul(four, scale)

    try:
        yield PiApproximation(0, _format_scaled(total, scale_power, digits))

        counter: Iterable[int]
        if iterations is None:
            counter = itertools.count(1)
        else:
            counter = range(1, iterations + 1)

        for idx in counter:
            two_n = _literal(2 * idx)
            two_n_plus_one = _literal(2 * idx + 1)
            two_n_plus_two = _literal(2 * idx + 2)

            product = mul(two_n, two_n_plus_one)
            denominator = mul(product, two_n_plus_two)
            term = div(numerator, denominator)

            updated = add(total, term) if idx % 2 else sub(total, term)

            total.free()
            total = updated

            yield PiApproximation(idx, _format_scaled(total, scale_power, digits))

            for number in (
                two_n,
                two_n_plus_one,
                two_n_plus_two,
                product,
                denominator,
                term,
            ):
                number.free()

    finally:
        for number in (total, numerator, scale, three, four):
            number.free()


def compute_pi_text(
    iterations: int,
    *,
    digits: int,
    extra_precision: int = 4,
) -> str:
    """Return the final Nilakantha approximation after the requested iterations."""
    approximation: str = "0"
    for state in iterate_nilakantha(
        iterations,
        digits=digits,
        extra_precision=extra_precision,
    ):
        approximation = state.digits
    return approximation