riichienv-core 0.3.2

#[cfg(feature = "python")]
use flate2::read::GzDecoder;
#[cfg(feature = "python")]
use pyo3::exceptions::PyValueError;
#[cfg(feature = "python")]
use pyo3::prelude::*;
use serde::{Deserialize, Serialize};
#[cfg(feature = "python")]
use std::fs::File;
#[cfg(feature = "python")]
use std::io::{BufReader, Read};
#[cfg(feature = "python")]
use std::sync::Arc;

#[cfg(feature = "python")]
use crate::replay::{Action, HuleData, LogKyoku, TileConverter, WinResultContextIterator};
#[cfg(feature = "python")]
use crate::types::MeldType;

#[cfg(feature = "python")]
#[pyclass(module = "riichienv._riichienv")]
pub struct MjSoulReplay {
    pub rounds: Vec<LogKyoku>,
}

#[cfg(feature = "python")]
#[derive(Debug)]
#[pyclass(module = "riichienv._riichienv")]
pub struct KyokuIterator {
    game: Py<MjSoulReplay>,
    index: usize,
    len: usize,
}

#[cfg(feature = "python")]
#[pymethods]
impl KyokuIterator {
    fn __iter__(slf: PyRef<'_, Self>) -> PyRef<'_, Self> {
        slf
    }

    fn __next__(mut slf: PyRefMut<'_, Self>) -> Option<LogKyoku> {
        if slf.index >= slf.len {
            return None;
        }

        let kyoku = {
            let game = slf.game.borrow(slf.py());
            game.rounds[slf.index].clone()
        };
        slf.index += 1;

        Some(kyoku)
    }
}

#[derive(Deserialize, Serialize, Clone, Debug)]
#[serde(tag = "name", content = "data")]
pub enum RawAction {
    #[serde(rename = "NewRound")]
    NewRound {
        scores: Vec<i32>,
        doras: Option<Vec<String>>,
        dora_indicators: Option<Vec<String>>,
        dora_marker: Option<String>,
        tiles0: Vec<String>,
        tiles1: Vec<String>,
        tiles2: Vec<String>,
        tiles3: Vec<String>,
        chang: u8,
        ju: u8,
        ben: Option<u8>,
        honba: Option<u8>,
        liqibang: u8,
        left_tile_count: Option<u8>,
        ura_doras: Option<Vec<String>>,
        paishan: Option<String>,
    },
    #[serde(rename = "DiscardTile")]
    DiscardTile {
        seat: usize,
        tile: String,
        #[serde(default)]
        is_liqi: bool,
        #[serde(default)]
        is_wliqi: bool,
        #[serde(default)]
        doras: Vec<String>,
    },
    #[serde(rename = "DealTile")]
    DealTile {
        seat: usize,
        tile: String,
        #[serde(default)]
        doras: Vec<String>,
        dora_marker: Option<String>,
        left_tile_count: Option<u8>,
    },
    #[serde(rename = "ChiPengGang")]
    ChiPengGang {
        seat: usize,
        #[serde(rename = "type")]
        meld_type: u64,
        tiles: Vec<String>,
        froms: Vec<usize>,
    },
    #[serde(rename = "AnGangAddGang")]
    AnGangAddGang {
        seat: usize,
        #[serde(rename = "type")]
        meld_type: u64,
        tiles: String,
    },
    #[serde(rename = "Hule")]
    Hule { hules: Vec<HuleDataRaw> },
    #[serde(rename = "dora")]
    Dora { dora_marker: String },
    #[serde(rename = "NoTile")]
    NoTile {},
    #[serde(rename = "BaBei")]
    BaBei {
        seat: usize,
        #[serde(default)]
        moqie: bool,
        #[serde(default)]
        doras: Vec<String>,
    },
    #[serde(rename = "LiuJu")]
    LiuJu {
        #[serde(rename = "type", default)]
        lj_type: u8,
        #[serde(default)]
        seat: usize,
        #[serde(default)]
        tiles: Vec<String>,
    },
    #[serde(other)]
    Other,
}

#[derive(Deserialize, Serialize, Clone, Debug)]
pub struct HuleDataRaw {
    pub seat: usize,
    pub hu_tile: String,
    pub zimo: bool,
    pub count: u32,
    pub fu: u32,
    pub fans: Vec<FanRaw>,
    pub hand: Vec<String>,
    pub ura_dora_indicators: Option<Vec<String>>,
    pub li_doras: Option<Vec<String>>,
    pub yiman: bool,
    pub point_rong: u32,
    pub point_zimo_qin: u32,
    pub point_zimo_xian: u32,
}

#[derive(Deserialize, Serialize, Clone, Debug)]
pub struct FanRaw {
    pub id: u32,
    #[serde(default)]
    pub val: u32,
}

#[derive(Deserialize, Serialize)]
pub struct GameLog {
    pub rounds: Vec<Vec<RawAction>>,
}

#[cfg(feature = "python")]
#[pymethods]
impl MjSoulReplay {
    #[staticmethod]
    fn from_json(path: String) -> PyResult<Self> {
        let file = File::open(&path)
            .map_err(|e| PyValueError::new_err(format!("Failed to open file: {}", e)))?;
        let reader = BufReader::with_capacity(65536, file);
        let mut decoder = GzDecoder::new(reader);
        let mut buffer = Vec::with_capacity(128 * 1024);

        decoder
            .read_to_end(&mut buffer)
            .map_err(|e| PyValueError::new_err(format!("Failed to decompress: {}", e)))?;

        let log: GameLog = serde_json::from_slice(&buffer)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse JSON: {}", e)))?;

        let mut rounds = Vec::with_capacity(log.rounds.len());
        for r_raw in log.rounds {
            rounds.push(Self::kyoku_from_raw_actions(r_raw));
        }

        // Populate end_scores based on next round's start scores
        for i in 0..rounds.len().saturating_sub(1) {
            rounds[i].end_scores = rounds[i + 1].scores.clone();
        }

        Ok(MjSoulReplay { rounds })
    }

    #[staticmethod]
    fn from_dict(py: Python, paifu: Py<PyAny>) -> PyResult<Self> {
        let json = py.import("json")?;
        let s: String = json.call_method1("dumps", (paifu,))?.extract()?;
        let v: serde_json::Value = serde_json::from_str(&s)
            .map_err(|e| PyValueError::new_err(format!("Failed to parse JSON: {}", e)))?;

        let (rounds_raw, _rule) = if let Some(obj) = v.as_object() {
            if let Some(data) = obj.get("data") {
                // assume Paifu struct { header, data }
                let rounds: Vec<Vec<RawAction>> = serde_json::from_value(data.clone())
                    .map_err(|e| PyValueError::new_err(format!("Failed to parse rounds: {}", e)))?;
                // TODO: Parse header for rule if converting from Paifu
                (rounds, crate::rule::GameRule::default_mjsoul())
            } else {
                // maybe just dict of rounds? Unlikely given usage.
                return Err(PyValueError::new_err("Invalid dict format: missing 'data'"));
            }
        } else if v.is_array() {
            let rounds: Vec<Vec<RawAction>> = serde_json::from_value(v).map_err(|e| {
                PyValueError::new_err(format!("Failed to parse rounds list: {}", e))
            })?;
            (rounds, crate::rule::GameRule::default_mjsoul())
        } else {
            return Err(PyValueError::new_err(
                "Invalid input format: expected dict or list",
            ));
        };

        // Detect 3P from the first round's scores length
        let is_3p = rounds_raw
            .first()
            .and_then(|r| {
                if let RawAction::NewRound { scores, .. } = &r[0] {
                    Some(scores.len() == 3)
                } else {
                    None
                }
            })
            .unwrap_or(false);
        let rule = if is_3p {
            crate::rule::GameRule::default_mjsoul_sanma()
        } else {
            _rule
        };

        let mut rounds = Vec::with_capacity(rounds_raw.len());
        for r_raw in rounds_raw {
            let mut kyoku = Self::kyoku_from_raw_actions(r_raw);
            kyoku.rule = rule;
            rounds.push(kyoku);
        }

        // Populate end_scores based on next round's start scores
        for i in 0..rounds.len().saturating_sub(1) {
            rounds[i].end_scores = rounds[i + 1].scores.clone();
        }

        // Calculate game end scores using the last round
        let is_3p = rounds.first().map(|r| r.scores.len() == 3).unwrap_or(false);

        let game_end_scores = if let Some(last) = rounds.last_mut() {
            if is_3p {
                // For 3P, simulate using GameState3P
                let mut state = crate::state_3p::GameState3P::new(0, false, None, 0, last.rule);
                let initial_scores: [i32; 3] = last.scores.clone().try_into().unwrap_or([35000; 3]);
                let oya = last.ju % 3;
                let bakaze = match last.chang {
                    0 => crate::types::Wind::East,
                    1 => crate::types::Wind::South,
                    2 => crate::types::Wind::West,
                    3 => crate::types::Wind::North,
                    _ => crate::types::Wind::East,
                } as u8;
                state._initialize_round(
                    oya,
                    bakaze,
                    last.ben,
                    last.liqibang as u32,
                    None,
                    Some(initial_scores.to_vec()),
                );
                // Replace the randomly-dealt hands with the actual replay
                // hands so that tenpai detection in NoTile is correct.
                for (i, hand) in last.hands.iter().enumerate() {
                    if i < state.players.len() {
                        state.players[i].hand = hand.clone();
                        state.players[i].hand.sort();
                    }
                }
                for action in last.actions.iter() {
                    state.apply_log_action(action);
                }
                last.end_scores = state.players.iter().map(|p| p.score).collect();
                Some(last.end_scores.clone())
            } else {
                // 4P path
                let mut state = crate::state::GameState::new(0, false, None, 0, last.rule);
                let initial_scores: [i32; 4] = last.scores.clone().try_into().unwrap_or([25000; 4]);
                let oya = last.ju % 4;
                let bakaze = match last.chang {
                    0 => crate::types::Wind::East,
                    1 => crate::types::Wind::South,
                    2 => crate::types::Wind::West,
                    3 => crate::types::Wind::North,
                    _ => crate::types::Wind::East,
                } as u8;
                state._initialize_round(
                    oya,
                    bakaze,
                    last.ben,
                    last.liqibang as u32,
                    None,
                    Some(initial_scores.to_vec()),
                );
                // Replace the randomly-dealt hands with the actual replay
                // hands so that tenpai detection in NoTile is correct.
                for (i, hand) in last.hands.iter().enumerate() {
                    if i < state.players.len() {
                        state.players[i].hand = hand.clone();
                        state.players[i].hand.sort();
                    }
                }
                for action in last.actions.iter() {
                    state.apply_log_action(action);
                }
                last.end_scores = state.players.iter().map(|p| p.score).collect();
                Some(last.end_scores.clone())
            }
        } else {
            None
        };

        // Set game_end_scores for all rounds
        if let Some(ges) = game_end_scores {
            for r in &mut rounds {
                r.game_end_scores = Some(ges.clone());
            }
        }

        Ok(MjSoulReplay { rounds })
    }

    fn num_rounds(&self) -> usize {
        self.rounds.len()
    }

    fn take_kyokus(slf: Py<Self>, py: Python<'_>) -> PyResult<KyokuIterator> {
        let logs_len = slf.borrow(py).rounds.len();
        Ok(KyokuIterator {
            game: slf,
            index: 0,
            len: logs_len,
        })
    }

    fn verify(&self) -> (usize, usize) {
        let mut total_agari = 0;
        let mut total_mismatches = 0;

        for kyoku in &self.rounds {
            let mut iter = WinResultContextIterator::new(kyoku.clone());

            while let Some(ctx) = iter.do_next() {
                total_agari += 1;

                let sim_han = ctx.actual.han;
                let sim_fu = ctx.actual.fu;
                let sim_yaku = ctx.actual.yaku.clone();

                let exp_han = ctx.expected_han;
                let exp_fu = ctx.expected_fu;
                let exp_yaku = ctx.expected_yaku.clone();

                // IGNORED: 31 (Dora), 32 (Aka), 33 (Ura)
                let ignored = [31, 32, 33];
                let yakuman_ids: Vec<u32> = (35..51).collect();

                let mut sim_filtered: Vec<u32> = sim_yaku
                    .iter()
                    .filter(|y| !ignored.contains(y))
                    .cloned()
                    .collect();
                let mut exp_filtered: Vec<u32> = exp_yaku
                    .iter()
                    .filter(|y| !ignored.contains(y))
                    .cloned()
                    .collect();

                let mut normalized_exp_han = exp_han;
                let is_yakuman = exp_yaku.iter().any(|y| yakuman_ids.contains(y));
                if is_yakuman && exp_han < 13 {
                    normalized_exp_han = exp_han * 13;
                }

                let mut mismatch = false;
                sim_filtered.sort();
                exp_filtered.sort();

                if sim_filtered != exp_filtered {
                    mismatch = true;
                } else {
                    let sim_ignored_han =
                        sim_yaku.iter().filter(|y| ignored.contains(y)).count() as u32;
                    let exp_ignored_han =
                        exp_yaku.iter().filter(|y| ignored.contains(y)).count() as u32;
                    let expected_sim_han =
                        normalized_exp_han as i32 - exp_ignored_han as i32 + sim_ignored_han as i32;

                    if normalized_exp_han < 13 && sim_han as i32 != expected_sim_han {
                        if sim_han != normalized_exp_han {
                            mismatch = true;
                        }
                    } else if (sim_han >= 13) != (normalized_exp_han >= 13) {
                        mismatch = true;
                    }

                    if !mismatch && normalized_exp_han < 13 && sim_fu != exp_fu {
                        mismatch = true;
                    }
                }

                if mismatch {
                    total_mismatches += 1;
                    println!(
                        "Mismatch: seat={}, han=(sim={}, exp={}), fu=(sim={}, exp={})",
                        ctx.seat, sim_han, exp_han, sim_fu, exp_fu
                    );
                    println!("  Expected Yaku: {:?}", exp_yaku);
                    println!("  Actual Yaku: {:?}", sim_yaku);
                    println!("  Conditions: {:?}", ctx.conditions);
                }
            }
        }
        (total_agari, total_mismatches)
    }
}

#[cfg(feature = "python")]
impl MjSoulReplay {
    fn kyoku_from_raw_actions(raw_actions: Vec<RawAction>) -> LogKyoku {
        let mut scores = Vec::new();
        let mut doras = Vec::new();
        let mut hands = vec![Vec::new(); 4];
        let mut chang = 0;
        let mut ju = 0;
        let mut ben = 0;
        let mut liqibang = 0;
        let mut left_tile_count = 70;
        let mut ura_doras = Vec::new();
        let mut paishan = None;

        if let RawAction::NewRound {
            scores: s,
            doras: d_opt,
            dora_indicators,
            dora_marker,
            tiles0,
            tiles1,
            tiles2,
            tiles3,
            chang: c,
            ju: j,
            ben: b,
            honba,
            liqibang: l,
            left_tile_count: lc,
            ura_doras: ud,
            paishan: p,
        } = &raw_actions[0]
        {
            scores = s.clone();
            if let Some(da) = dora_indicators.as_ref().or(d_opt.as_ref()) {
                for v in da {
                    doras.push(TileConverter::parse_tile_136(v));
                }
            } else if let Some(dm) = dora_marker {
                doras.push(TileConverter::parse_tile_136(dm));
            }
            hands = vec![
                tiles0
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect(),
                tiles1
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect(),
                tiles2
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect(),
                tiles3
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect(),
            ];
            chang = *c;
            ju = *j;
            ben = b.or(*honba).unwrap_or(0);
            liqibang = *l;
            left_tile_count = lc.unwrap_or(70);
            if let Some(uda) = ud {
                ura_doras = uda
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect();
            }
            paishan = p.clone();
        }

        let mut actions = Vec::with_capacity(raw_actions.len());
        for ma in raw_actions {
            actions.push(Self::parse_raw_action(ma));
        }

        let end_scores = scores.clone();

        let mut wliqi = vec![false; 4];
        for action in &actions {
            if let Action::DiscardTile { seat, is_wliqi, .. } = action {
                if *is_wliqi {
                    wliqi[*seat] = true;
                }
            }
        }

        LogKyoku {
            scores,
            end_scores,
            doras,
            ura_doras,
            hands,
            chang,
            ju,
            ben,
            liqibang,
            left_tile_count,
            wliqi,
            paishan,
            actions: Arc::from(actions),
            rule: crate::rule::GameRule::default_mjsoul(),
            game_end_scores: None,
        }
    }

    fn parse_raw_action(ma: RawAction) -> Action {
        match ma {
            RawAction::DiscardTile {
                seat,
                tile,
                is_liqi,
                is_wliqi,
                doras,
            } => Action::DiscardTile {
                seat,
                tile: TileConverter::parse_tile_136(&tile),
                is_liqi,
                is_wliqi,
                doras: if doras.is_empty() {
                    None
                } else {
                    Some(
                        doras
                            .iter()
                            .map(|v| TileConverter::parse_tile_136(v))
                            .collect(),
                    )
                },
            },
            RawAction::DealTile {
                seat,
                tile,
                doras,
                dora_marker,
                left_tile_count,
            } => {
                let mut d_res = if doras.is_empty() {
                    None
                } else {
                    Some(
                        doras
                            .iter()
                            .map(|v| TileConverter::parse_tile_136(v))
                            .collect(),
                    )
                };
                if d_res.is_none() {
                    if let Some(dm) = dora_marker {
                        d_res = Some(vec![TileConverter::parse_tile_136(&dm)]);
                    }
                }
                Action::DealTile {
                    seat,
                    tile: TileConverter::parse_tile_136(&tile),
                    doras: d_res,
                    left_tile_count,
                }
            }
            RawAction::ChiPengGang {
                seat,
                meld_type,
                tiles,
                froms,
            } => {
                let m_type = match meld_type {
                    0 => MeldType::Chi,
                    1 => MeldType::Pon,
                    2 => MeldType::Daiminkan,
                    3 => MeldType::Ankan,
                    _ => MeldType::Chi,
                };
                Action::ChiPengGang {
                    seat,
                    meld_type: m_type,
                    tiles: tiles
                        .iter()
                        .map(|v| TileConverter::parse_tile_136(v))
                        .collect(),
                    froms,
                }
            }
            RawAction::AnGangAddGang {
                seat,
                meld_type,
                tiles,
            } => {
                let m_type = if meld_type == 3 {
                    MeldType::Ankan
                } else {
                    MeldType::Kakan
                };
                let tile_raw_id = TileConverter::parse_tile_34(&tiles).0;
                Action::AnGangAddGang {
                    seat,
                    meld_type: m_type,
                    tiles: vec![TileConverter::parse_tile_136(&tiles)],
                    tile_raw_id,
                    doras: None, // Will be updated by Dora action or DealTile
                }
            }
            RawAction::Hule { hules } => {
                let hules_typed = hules
                    .into_iter()
                    .map(|h| HuleData {
                        seat: h.seat,
                        hu_tile: TileConverter::parse_tile_136(&h.hu_tile),
                        zimo: h.zimo,
                        count: h.count,
                        fu: h.fu,
                        fans: h.fans.iter().filter(|f| f.val > 0).map(|f| f.id).collect(),
                        li_doras: h
                            .ura_dora_indicators
                            .or(h.li_doras)
                            .map(|a| a.iter().map(|v| TileConverter::parse_tile_136(v)).collect()),
                        yiman: h.yiman,
                        point_rong: h.point_rong,
                        point_zimo_qin: h.point_zimo_qin,
                        point_zimo_xian: h.point_zimo_xian,
                    })
                    .collect();
                Action::Hule { hules: hules_typed }
            }
            RawAction::Dora { dora_marker } => Action::Dora {
                dora_marker: TileConverter::parse_tile_136(&dora_marker),
            },
            RawAction::NoTile {} => Action::NoTile,
            RawAction::BaBei {
                seat,
                moqie,
                doras: _,
            } => Action::BaBei { seat, moqie },
            RawAction::LiuJu {
                lj_type,
                seat,
                tiles,
            } => Action::LiuJu {
                lj_type,
                seat,
                tiles: tiles
                    .iter()
                    .map(|v| TileConverter::parse_tile_136(v))
                    .collect(),
            },
            _ => Action::Other("Other".to_string()),
        }
    }
}