use alloc::collections::VecDeque;

use rustpython_common::wtf8::{Wtf8, Wtf8Buf, wtf8_concat};

use std::collections::HashMap;

use crate::_datetime::{PyTzInfo, datetime_type, timedelta_from_seconds};

use crate::common::lock::PyMutex;

use crate::vm::{

AsObject, Py, PyObject, PyObjectRef, PyPayload, PyRef, PyResult, TryFromObject,

VirtualMachine,

builtins::{

PyBaseExceptionRef, PyBytes, PyModule, PyStr, PyStrRef, PyTuple, PyType, PyTypeRef,

},

class::StaticType,

function::{ArgIndex, FuncArgs, IntoFuncArgs, KwArgs},

object::{Traverse, TraverseFn},

protocol::{PyIter, PyIterReturn},

types::{Constructor, Representable},

};

const EPOCHORDINAL: i64 = 719_163;

const SOURCE_NOCACHE: u8 = 0;

const SOURCE_CACHE: u8 = 1;

const SOURCE_FILE: u8 = 2;

const STRONG_CACHE_MAX: usize = 8;

const DELTA_CACHE_MAX: usize = 512;

const DAYS_IN_MONTH: [i32; 13] = [-1, 31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31];

const DAYS_BEFORE_MONTH: [i32; 13] =

[-1, 0, 31, 59, 90, 120, 151, 181, 212, 243, 273, 304, 334];

#[derive(Clone, Copy, Debug)]

enum DateRule {

Calendar {

month: u8,

week: u8,

day: u8,

hour: i32,

minute: i32,

second: i32,

},

Day {

julian: bool,

day: u16,

hour: i32,

minute: i32,

second: i32,

},

}

impl DateRule {

fn timestamp(self, year: i32) -> i64 {

match self {

Self::Calendar {

month,

week,

day,

hour,

minute,

second,

} => calendar_timestamp(year, month, week, day, hour, minute, second),

Self::Day {

julian,

day,

hour,

minute,

second,

} => day_timestamp(year, julian, day, hour, minute, second),

}

}

}

#[derive(Clone, Debug, Traverse)]

struct TtInfo {

utcoff: PyObjectRef,

dstoff: PyObjectRef,

tzname: PyObjectRef,

#[pytraverse(skip)]

utcoff_seconds: i64,

}

#[derive(Debug, Traverse)]

struct TzRule {

std: TtInfo,

dst: Option<TtInfo>,

#[pytraverse(skip)]

std_only: bool,

#[pytraverse(skip)]

dst_diff: i64,

#[pytraverse(skip)]

start: Option<DateRule>,

#[pytraverse(skip)]

end: Option<DateRule>,

}

#[derive(Debug, Traverse)]

struct ZoneData {

#[pytraverse(skip)]

trans_utc: Vec<i64>,

// fold 0 keeps the larger offset, fold 1 the smaller.

#[pytraverse(skip)]

trans_wall: [Vec<i64>; 2],

ttinfos: Vec<TtInfo>,

#[pytraverse(skip)]

trans_ttinfos: Vec<usize>,

#[pytraverse(skip)]

tti_before: Option<usize>,

after: TzRule,

#[pytraverse(skip)]

fixed_offset: bool,

}

// The i64 fields would raise this value's alignment above a pointer.

// On 32-bit that moves the payload inside `Py<T>` relative to `PyTzInfo`.

#[derive(Debug)]

struct ZoneDataBox(alloc::boxed::Box<ZoneData>);

impl core::ops::Deref for ZoneDataBox {

type Target = ZoneData;

fn deref(&self) -> &Self::Target {

&self.0

}

}

unsafe impl Traverse for ZoneDataBox {

fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {

self.0.traverse(traverse_fn);

}

}

#[derive(Clone, Copy, Debug)]

enum TtSel {

Index(usize),

Std,

Dst,

None,

}

enum TzFail {

StdFormat,

StdOffset,

DstFormat,

DstOffset,

MissingRules,

MalformedRule,

Extra,

}

#[derive(Debug)]

struct CacheState {

strong: VecDeque<(PyObjectRef, PyObjectRef)>,

weak: Option<PyObjectRef>,

deltas: HashMap<i64, PyObjectRef>,

delta_order: VecDeque<i64>,

}

fn empty_cache() -> CacheState {

CacheState {

strong: VecDeque::new(),

weak: None,

deltas: HashMap::new(),

delta_order: VecDeque::new(),

}

}

unsafe impl Traverse for CacheState {

fn traverse(&self, traverse_fn: &mut TraverseFn<'_>) {

for (key, zone) in &self.strong {

key.traverse(traverse_fn);

zone.traverse(traverse_fn);

}

self.weak.traverse(traverse_fn);

for seconds in &self.delta_order {

let Some(delta) = self.deltas.get(seconds) else {

continue;

};

delta.traverse(traverse_fn);

}

}

}

// One of these per imported module. A subinterpreter has its own module,

// so its cache is not the parent's and drops with that interpreter.

#[pyclass(no_attr, module = "_zoneinfo", name = "_zoneinfo_state", traverse)]

#[derive(Debug, PyPayload)]

struct ZoneInfoState {

cache: PyMutex<CacheState>,

}

#[pyclass(flags(DISALLOW_INSTANTIATION))]

impl ZoneInfoState {}

fn module_state(vm: &VirtualMachine) -> PyResult<PyRef<ZoneInfoState>> {

let modules = vm.sys_module.get_attr("modules", vm)?;

let module = match modules.get_item("_zoneinfo", vm) {

Ok(module) => module,

Err(_) => vm.import("_zoneinfo", 0)?,

};

module

.get_attr("_state", vm)?

.downcast::<ZoneInfoState>()

.map_err(|_| vm.new_runtime_error("_zoneinfo state is corrupted"))

}

fn with_cache<R>(vm: &VirtualMachine, f: impl FnOnce(&mut CacheState) -> R) -> PyResult<R> {

let state = module_state(vm)?;

let mut cache = state.cache.lock();

Ok(f(&mut cache))

}

fn is_base(cls: &Py<PyType>) -> bool {

cls.is(ZoneInfo::static_type())

}

fn strong_entries(vm: &VirtualMachine) -> PyResult<Vec<(PyObjectRef, PyObjectRef)>> {

with_cache(vm, |cache| cache.strong.iter().cloned().collect())

}

fn move_strong_to_front(zone: &PyObject, vm: &VirtualMachine) -> PyResult<()> {

with_cache(vm, |cache| {

let Some(pos) = cache.strong.iter().position(|(_, cached)| cached.is(zone)) else {

return;

};

if pos != 0

&& let Some(entry) = cache.strong.remove(pos)

{

cache.strong.push_front(entry);

}

})

}

fn strong_cache_get(

cls: &Py<PyType>,

key: &PyObject,

vm: &VirtualMachine,

) -> PyResult<Option<PyObjectRef>> {

if !is_base(cls) {

return Ok(None);

}

for (cached_key, zone) in strong_entries(vm)? {

if vm.bool_eq(key, &cached_key)? {

move_strong_to_front(&zone, vm)?;

return Ok(Some(zone));

}

}

Ok(None)

}

fn strong_cache_put(

cls: &Py<PyType>,

key: PyObjectRef,

zone: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<()> {

if !is_base(cls) {

return Ok(());

}

let evicted = with_cache(vm, |cache| {

cache.strong.push_front((key, zone));

cache

.strong

.split_off(STRONG_CACHE_MAX.min(cache.strong.len()))

})?;

// Dropping a zone or its key can reenter the cache through Python.

drop(evicted);

Ok(())

}

fn strong_cache_eject(cls: &Py<PyType>, key: &PyObject, vm: &VirtualMachine) -> PyResult<()> {

if !is_base(cls) {

return Ok(());

}

for (cached_key, zone) in strong_entries(vm)? {

if vm.bool_eq(key, &cached_key)? {

let removed = with_cache(vm, |cache| {

cache

.strong

.iter()

.position(|(_, cached)| cached.is(&zone))

.and_then(|pos| cache.strong.remove(pos))

})?;

drop(removed);

return Ok(());

}

}

Ok(())

}

fn strong_cache_clear(cls: &Py<PyType>, vm: &VirtualMachine) -> PyResult<()> {

if !is_base(cls) {

return Ok(());

}

let removed = with_cache(vm, |cache| core::mem::take(&mut cache.strong))?;

drop(removed);

Ok(())

}

fn cached_delta(seconds: i64, vm: &VirtualMachine) -> PyResult<PyObjectRef> {

if let Some(delta) = with_cache(vm, |cache| cache.deltas.get(&seconds).cloned())? {

return Ok(delta);

}

let delta = timedelta_from_seconds(seconds, vm)?;

with_cache(vm, |cache| {

if let Some(existing) = cache.deltas.get(&seconds) {

return existing.clone();

}

if cache.deltas.len() >= DELTA_CACHE_MAX

&& let Some(old) = cache.delta_order.pop_front()

{

cache.deltas.remove(&old);

}

cache.deltas.insert(seconds, delta.clone());

cache.delta_order.push_back(seconds);

delta

})

}

fn new_weak_cache(vm: &VirtualMachine) -> PyResult {

let weakref = vm.import("weakref", 0)?;

let dict = weakref.get_attr("WeakValueDictionary", vm)?;

dict.call((), vm)

}

fn base_weak_cache(vm: &VirtualMachine) -> PyResult {

if let Some(weak) = with_cache(vm, |cache| cache.weak.clone())? {

return Ok(weak);

}

let weak = new_weak_cache(vm)?;

with_cache(vm, |cache| {

if let Some(existing) = &cache.weak {

return existing.clone();

}

cache.weak = Some(weak.clone());

weak.clone()

})

}

fn get_weak_cache(cls: &Py<PyType>, vm: &VirtualMachine) -> PyResult {

if is_base(cls) {

base_weak_cache(vm)

} else {

cls.as_object().get_attr("_weak_cache", vm)

}

}

fn call_attr(

module: &'static str,

name: &'static str,

args: impl IntoFuncArgs,

vm: &VirtualMachine,

) -> PyResult {

// import() returns the top-level package for a dotted name.

let imported = vm.import(module, 0)?;

let module = match module.rsplit_once('.') {

Some((_, attr)) => imported.get_attr(attr, vm)?,

None => imported,

};

let func = module.get_attr(name, vm)?;

func.call(args, vm)

}

fn floor_mod(value: i64, modulo: i64) -> i64 {

let mut out = value % modulo;

if out < 0 {

out += modulo;

}

out

}

fn is_leap_year(year: i32) -> bool {

let year = year as u32;

year.is_multiple_of(4) && (!year.is_multiple_of(100) || year.is_multiple_of(400))

}

fn ymd_to_ord(y: i32, m: i32, d: i32) -> i32 {

let y = y - 1;

let days_before_year = y * 365 + y / 4 - y / 100 + y / 400;

let mut yearday = DAYS_BEFORE_MONTH[m as usize];

if m > 2 && is_leap_year(y + 1) {

yearday += 1;

}

days_before_year + yearday + d

}

fn calendar_timestamp(

year: i32,

month: u8,

week: u8,

day: u8,

hour: i32,

minute: i32,

second: i32,

) -> i64 {

let month = i32::from(month);

let first_day = floor_mod(i64::from(ymd_to_ord(year, month, 1)) + 6, 7);

let mut days_in_month = DAYS_IN_MONTH[month as usize];

if month == 2 && is_leap_year(year) {

days_in_month += 1;

}

let mut month_day = floor_mod(i64::from(day) - (first_day + 1), 7) + 1;

month_day += (i64::from(week) - 1) * 7;

if month_day > i64::from(days_in_month) {

month_day -= 7;

}

let ordinal = i64::from(ymd_to_ord(year, month, month_day as i32)) - EPOCHORDINAL;

ordinal * 86_400 + i64::from(hour) * 3_600 + i64::from(minute) * 60 + i64::from(second)

}

fn day_timestamp(

year: i32,

julian: bool,

day: u16,

hour: i32,

minute: i32,

second: i32,

) -> i64 {

let days_before_year = i64::from(ymd_to_ord(year, 1, 1)) - EPOCHORDINAL - 1;

let mut day = i64::from(day);

if julian && day >= 59 && is_leap_year(year) {

day += 1;

}

(days_before_year + day) * 86_400

+ i64::from(hour) * 3_600

+ i64::from(minute) * 60

+ i64::from(second)

}

fn bisect_right(value: i64, arr: &[i64]) -> usize {

arr.partition_point(|&item| item <= value)

}

fn order_offsets(a: i64, b: i64) -> (i64, i64) {

if b > a { (b, a) } else { (a, b) }

}

fn utcoff_to_dstoff(trans_idx: &[usize], utcoffs: &[i64], isdsts: &[bool]) -> Vec<i64> {

let num_ttinfos = utcoffs.len();

let num_transitions = trans_idx.len();

let mut dstoffs = vec![0_i64; num_ttinfos];

let dst_count = num_ttinfos;

let mut dst_found = 0_usize;

for i in 1..num_transitions {

if dst_count == dst_found {

break;

}

let idx = trans_idx[i];

let mut comp_idx = trans_idx[i - 1];

if !isdsts[idx] || dstoffs[idx] != 0 {

continue;

}

let utcoff = utcoffs[idx];

let mut dstoff = if !isdsts[comp_idx] {

utcoff - utcoffs[comp_idx]

} else {

0

};

if dstoff == 0 && idx + 1 < num_ttinfos && i + 1 < num_transitions {

comp_idx = trans_idx[i + 1];

if isdsts[comp_idx] {

continue;

}

dstoff = utcoff - utcoffs[comp_idx];

}

if dstoff != 0 {

dst_found += 1;

dstoffs[idx] = dstoff;

}

}

if dst_found < dst_count {

for idx in 0..num_ttinfos {

if isdsts[idx] && dstoffs[idx] == 0 {

dstoffs[idx] = 3_600;

}

}

}

dstoffs

}

fn ts_to_local(trans_idx: &[usize], trans_utc: &[i64], utcoff: &[i64]) -> [Vec<i64>; 2] {

let n = trans_utc.len();

if n == 0 {

return [Vec::new(), Vec::new()];

}

let mut wall0 = trans_utc.to_vec();

let mut wall1 = trans_utc.to_vec();

let (offset_0, offset_1) = if utcoff.len() > 1 {

order_offsets(utcoff[0], utcoff[trans_idx[0]])

} else {

(utcoff[0], utcoff[0])

};

wall0[0] += offset_0;

wall1[0] += offset_1;

for i in 1..n {

let (off0, off1) = order_offsets(utcoff[trans_idx[i - 1]], utcoff[trans_idx[i]]);

wall0[i] += off0;

wall1[i] += off1;

}

[wall0, wall1]

}

struct Parser<'a> {

bytes: &'a [u8],

i: usize,

}

impl Parser<'_> {

fn peek(&self) -> Option<u8> {

match self.bytes.get(self.i).copied() {

Some(0) | None => None,

Some(byte) => Some(byte),

}

}

fn bump(&mut self) -> Option<u8> {

let byte = self.peek()?;

self.i += 1;

Some(byte)

}

}

fn parse_digits(parser: &mut Parser<'_>, min: usize, max: usize) -> Option<i32> {

let mut value = 0_i32;

let mut count = 0_usize;

while count < max {

let Some(byte) = parser.peek() else {

return (count >= min).then_some(value);

};

if !byte.is_ascii_digit() {

return (count >= min).then_some(value);

}

parser.bump();

value = value * 10 + i32::from(byte - b'0');

count += 1;

}

Some(value)

}

fn parse_abbr(parser: &mut Parser<'_>) -> Option<String> {

if parser.peek() == Some(b'<') {

parser.bump();

let start = parser.i;

loop {

let byte = parser.peek()?;

if byte == b'>' {

break;

}

if !(byte.is_ascii_alphanumeric() || byte == b'+' || byte == b'-') {

return None;

}

parser.bump();

}

let end = parser.i;

if end == start {

return None;

}

parser.bump();

return Some(String::from_utf8_lossy(&parser.bytes[start..end]).into_owned());

}

let start = parser.i;

while parser.peek().is_some_and(|byte| byte.is_ascii_alphabetic()) {

parser.bump();

}

if parser.i == start {

return None;

}

Some(String::from_utf8_lossy(&parser.bytes[start..parser.i]).into_owned())

}

fn parse_transition_time(parser: &mut Parser<'_>) -> Option<(i32, i32, i32)> {

let mut sign = 1_i32;

if matches!(parser.peek(), Some(b'-' | b'+')) {

if parser.peek() == Some(b'-') {

sign = -1;

}

parser.bump();

}

let hour = parse_digits(parser, 1, 3)? * sign;

let mut minute = 0_i32;

let mut second = 0_i32;

if parser.peek() == Some(b':') {

parser.bump();

minute = parse_digits(parser, 2, 2)? * sign;

if parser.peek() == Some(b':') {

parser.bump();

second = parse_digits(parser, 2, 2)? * sign;

}

}

Some((hour, minute, second))

}

fn parse_tz_delta(parser: &mut Parser<'_>) -> Option<i64> {

let (hours, minutes, seconds) = parse_transition_time(parser)?;

if !(-24..=24).contains(&hours) {

return None;

}

Some(-((i64::from(hours) * 3_600) + (i64::from(minutes) * 60) + i64::from(seconds)))

}

fn calendar_rule(

month: i32,

week: i32,

day: i32,

hour: i32,

minute: i32,

second: i32,

) -> Option<DateRule> {

if !(1..=12).contains(&month)

|| !(1..=5).contains(&week)

|| !(0..=6).contains(&day)

|| !(-167..=167).contains(&hour)

{

return None;

}

Some(DateRule::Calendar {

month: month as u8,

week: week as u8,

day: day as u8,

hour,

minute,

second,

})

}

fn day_rule(julian: bool, day: i32, hour: i32, minute: i32, second: i32) -> Option<DateRule> {

let min_day = i32::from(julian);

if day < min_day || day > 365 || !(-167..=167).contains(&hour) {

return None;

}

Some(DateRule::Day {

julian,

day: day as u16,

hour,

minute,

second,

})

}

fn parse_transition_rule(parser: &mut Parser<'_>) -> Option<DateRule> {

let mut hour = 2_i32;

let mut minute = 0_i32;

let mut second = 0_i32;

if parser.peek() == Some(b'M') {

parser.bump();

let month = parse_digits(parser, 1, 2)?;

if parser.bump() != Some(b'.') {

return None;

}

let week = parse_digits(parser, 1, 1)?;

if parser.bump() != Some(b'.') {

return None;

}

let day = parse_digits(parser, 1, 1)?;

if parser.peek() == Some(b'/') {

parser.bump();

(hour, minute, second) = parse_transition_time(parser)?;

}

return calendar_rule(month, week, day, hour, minute, second);

}

let mut julian = false;

if parser.peek() == Some(b'J') {

julian = true;

parser.bump();

}

let day = parse_digits(parser, 1, 3)?;

if parser.peek() == Some(b'/') {

parser.bump();

(hour, minute, second) = parse_transition_time(parser)?;

}

day_rule(julian, day, hour, minute, second)

}

fn tz_error(kind: TzFail, shown: &str, vm: &VirtualMachine) -> PyBaseExceptionRef {

let message = match kind {

TzFail::StdFormat => format!("Invalid STD format in {shown}"),

TzFail::StdOffset => format!("Invalid STD offset in {shown}"),

TzFail::DstFormat => format!("Invalid DST format in {shown}"),

TzFail::DstOffset => format!("Invalid DST offset in {shown}"),

TzFail::MissingRules => format!("Missing transition rules in TZ string: {shown}"),

TzFail::MalformedRule => format!("Malformed transition rule in TZ string: {shown}"),

TzFail::Extra => format!("Extraneous characters at end of TZ string: {shown}"),

};

vm.new_value_error(message)

}

fn build_ttinfo(

utcoff_seconds: i64,

dst_seconds: i64,

tzname: PyObjectRef,

vm: &VirtualMachine,

) -> PyResult<TtInfo> {

Ok(TtInfo {

utcoff: cached_delta(utcoff_seconds, vm)?,

dstoff: cached_delta(dst_seconds, vm)?,

tzname,

utcoff_seconds,

})

}

fn build_tzrule(

std_abbr: PyObjectRef,

dst_abbr: Option<PyObjectRef>,

std_offset: i64,

dst_offset: i64,

start: Option<DateRule>,

end: Option<DateRule>,

vm: &VirtualMachine,

) -> PyResult<TzRule> {

let std = build_ttinfo(std_offset, 0, std_abbr, vm)?;

let (dst, std_only, dst_diff) = if let Some(dst_abbr) = dst_abbr {

let dst_diff = dst_offset - std_offset;

let dst = build_ttinfo(dst_offset, dst_diff, dst_abbr, vm)?;

(Some(dst), false, dst_diff)

} else {

(None, true, 0)

};

Ok(TzRule {

std,

dst,

std_only,

dst_diff,

start,

end,

})

}

fn parse_tz_str(bytes: &[u8], shown: &str, vm: &VirtualMachine) -> PyResult<TzRule> {

let mut parser = Parser { bytes, i: 0 };

let std_abbr =

parse_abbr(&mut parser).ok_or_else(|| tz_error(TzFail::StdFormat, shown, vm))?;

let std_offset =

parse_tz_delta(&mut parser).ok_or_else(|| tz_error(TzFail::StdOffset, shown, vm))?;

if parser.peek().is_none() {

return build_tzrule(

vm.ctx.new_str(std_abbr).into(),

None,

std_offset,

0,

None,

None,

vm,

);

}

let dst_abbr =

parse_abbr(&mut parser).ok_or_else(|| tz_error(TzFail::DstFormat, shown, vm))?;

let dst_offset = if parser.peek() == Some(b',') {

std_offset + 3_600

} else {

parse_tz_delta(&mut parser).ok_or_else(|| tz_error(TzFail::DstOffset, shown, vm))?

};

let mut rules = [None, None];

for rule in &mut rules {

if parser.peek() != Some(b',') {

return Err(tz_error(TzFail::MissingRules, shown, vm));

}

parser.bump();

*rule = Some(

parse_transition_rule(&mut parser)

.ok_or_else(|| tz_error(TzFail::MalformedRule, shown, vm))?,

);

}

if parser.peek().is_some() {

return Err(tz_error(TzFail::Extra, shown, vm));

}

build_tzrule(

vm.ctx.new_str(std_abbr).into(),

Some(vm.ctx.new_str(dst_abbr).into()),

std_offset,

dst_offset,

rules[0],

rules[1],

vm,

)

}

fn tuple_owned(tup: &Py<PyTuple>, index: usize, vm: &VirtualMachine) -> PyResult<PyObjectRef> {

tup.as_slice()

.get(index)

.cloned()

.ok_or_else(|| vm.new_index_error("tuple index out of range"))

}

fn tuple_i64(tup: &Py<PyTuple>, index: usize, vm: &VirtualMachine) -> PyResult<i64> {

tuple_owned(tup, index, vm)?.try_to_value(vm)

}

fn as_exact_tuple<'a>(obj: &'a PyObject, vm: &VirtualMachine) -> PyResult<&'a Py<PyTuple>> {

if let Some(tup) = obj.downcast_ref_if_exact::<PyTuple>(vm) {