"""
p99 — Design Log Storage System (LC 635)

INVARIANT (prefix range = chronological range): timestamps are fixed-width
  zero-padded "YYYY:MM:DD:HH:MM:SS"; therefore lexicographic comparison
  matches chronological comparison. To range-query at a granularity, truncate
  start/end to the granularity's prefix length.
"""
from __future__ import annotations

import bisect
import random
from typing import List, Tuple


_PREFIX_LEN = {
    "Year": 4,
    "Month": 7,
    "Day": 10,
    "Hour": 13,
    "Minute": 16,
    "Second": 19,
}


class LogSystem:
    """Linear scan with prefix-length dict. put O(1), retrieve O(n)."""

    def __init__(self) -> None:
        self.logs: List[Tuple[str, int]] = []

    def put(self, id: int, timestamp: str) -> None:
        self.logs.append((timestamp, id))

    def retrieve(self, start: str, end: str, granularity: str) -> List[int]:
        n = _PREFIX_LEN[granularity]
        s, e = start[:n], end[:n]
        return [i for (ts, i) in self.logs if s <= ts[:n] <= e]


class LogSystemSorted:
    """Sorted-list + bisect oracle. retrieve O(log n + k)."""

    def __init__(self) -> None:
        # Keep sorted by full timestamp (lex == chronological).
        self.timestamps: List[str] = []
        self.ids: List[int] = []

    def put(self, id: int, timestamp: str) -> None:
        idx = bisect.bisect_left(self.timestamps, timestamp)
        self.timestamps.insert(idx, timestamp)
        self.ids.insert(idx, id)

    def retrieve(self, start: str, end: str, granularity: str) -> List[int]:
        n = _PREFIX_LEN[granularity]
        s = start[:n]
        e = end[:n]
        # Lower bound on full timestamps: any ts with ts[:n] >= s.
        # Equivalent: ts >= s + "0..." padding to position n. Use a simple
        # padded floor and ceiling for clarity.
        # Use the smallest possible suffix ("0:0:0:0:0" style) for floor,
        # and the largest ("9:99:99:99:99") for ceiling — but it's easier
        # to just scan candidates near the bisect; for safety we fall back
        # to a filter against bisect bounds.
        lo = bisect.bisect_left(self.timestamps, s)
        # Anything <= e at prefix n is included.
        # Build a ceiling key by padding e with a char larger than any that
        # appears in a timestamp. The full format max is "9999:99:99:99:99:99";
        # ':' (0x3A) is the largest char used, so pad with '~' (0x7E) > ':'.
        ceil_key = e + "~" * (19 - n)
        hi = bisect.bisect_right(self.timestamps, ceil_key)
        return [self.ids[i] for i in range(lo, hi) if s <= self.timestamps[i][:n] <= e]


def edge_cases() -> None:
    ls = LogSystem()
    ls.put(1, "2017:01:01:23:59:59")
    ls.put(2, "2017:01:01:22:59:59")
    ls.put(3, "2016:01:01:00:00:00")
    r = ls.retrieve("2016:01:01:01:01:01", "2017:01:01:23:00:00", "Year")
    assert sorted(r) == [1, 2, 3], r
    # Hour granularity: start prefix "2016:01:01:01"; id=3 has hour prefix
    # "2016:01:01:00" which is < start, so it's excluded.
    r = ls.retrieve("2016:01:01:01:01:01", "2017:01:01:23:00:00", "Hour")
    assert sorted(r) == [1, 2], r
    r = ls.retrieve("2017:01:01:23:59:59", "2017:01:01:23:59:59", "Second")
    assert r == [1], r

    # Empty.
    ls2 = LogSystem()
    assert ls2.retrieve("2010:01:01:00:00:00", "2020:01:01:00:00:00", "Year") == []


def _random_ts(rng: random.Random) -> str:
    return (
        f"{rng.randint(2000, 2020):04d}:"
        f"{rng.randint(1, 12):02d}:"
        f"{rng.randint(1, 28):02d}:"
        f"{rng.randint(0, 23):02d}:"
        f"{rng.randint(0, 59):02d}:"
        f"{rng.randint(0, 59):02d}"
    )


def stress_test() -> None:
    rng = random.Random(42)
    grans = list(_PREFIX_LEN.keys())
    for _ in range(100):
        a = LogSystem()
        b = LogSystemSorted()
        n_puts = rng.randint(1, 30)
        for i in range(n_puts):
            ts = _random_ts(rng)
            a.put(i, ts)
            b.put(i, ts)
        for _ in range(5):
            s = _random_ts(rng)
            e = _random_ts(rng)
            if s > e:
                s, e = e, s
            g = rng.choice(grans)
            ra = sorted(a.retrieve(s, e, g))
            rb = sorted(b.retrieve(s, e, g))
            assert ra == rb, (s, e, g, ra, rb)


if __name__ == "__main__":
    edge_cases()
    stress_test()
    print("ALL TESTS PASSED")
