use snforge_std::{declare, ContractClassTrait, start_prank, stop_prank, CheatTarget};
use pyth::wormhole::{IWormholeDispatcher, IWormholeDispatcherTrait, ParseAndVerifyVmError};
use pyth::reader::ReaderImpl;
use pyth::byte_array::{ByteArray, ByteArrayImpl};
use pyth::util::{UnwrapWithFelt252, array_try_into};
use core::starknet::{ContractAddress, EthAddress};
use core::panic_with_felt252;

#[test]
fn test_parse_and_verify_vm_works() {
    let dispatcher = deploy_and_init();

    let vm = dispatcher.parse_and_verify_vm(good_vm1());
    assert!(vm.version == 1);
    assert!(vm.guardian_set_index == 3);
    assert!(vm.signatures.len() == 13);
    assert!(*vm.signatures.at(0).guardian_index == 1);
    assert!(*vm.signatures.at(1).guardian_index == 2);
    assert!(*vm.signatures.at(12).guardian_index == 18);
    assert!(vm.timestamp == 1712589207);
    assert!(vm.nonce == 0);
    assert!(vm.emitter_chain_id == 26);
    assert!(
        vm.emitter_address == 0xe101faedac5851e32b9b23b5f9411a8c2bac4aae3ed4dd7b811dd1a72ea4aa71
    );
    assert!(vm.sequence == 0x2f03161);
    assert!(vm.consistency_level == 1);
    assert!(vm.payload.len() == 37);

    let mut reader = ReaderImpl::new(vm.payload);
    assert!(reader.read_u8() == 65);
    assert!(reader.read_u8() == 85);
    assert!(reader.read_u8() == 87);
}

#[test]
#[fuzzer(runs: 100, seed: 0)]
#[should_panic]
fn test_parse_and_verify_vm_rejects_corrupted_vm(pos: usize, random1: usize, random2: usize) {
    let dispatcher = deploy_and_init();

    let input = corrupted_vm(good_vm1(), pos, random1, random2);
    let vm = dispatcher.parse_and_verify_vm(input);
    println!("no error, output: {:?}", vm);
}

#[test]
#[should_panic(expected: ('wrong governance contract',))]
fn test_submit_guardian_set_rejects_invalid_emitter() {
    let dispatcher = deploy(
        array_try_into(array![0x686b9ea8e3237110eaaba1f1b7467559a3273819]),
        CHAIN_ID,
        GOVERNANCE_CHAIN_ID,
        GOVERNANCE_CONTRACT
    );

    dispatcher
        .submit_new_guardian_set(
            ByteArrayImpl::new(
                array_try_into(
                    array![
                        1766847064779993780836504669527478016540696065743262550388106844126893447,
                        69151420679589009790918040295649622139220704596068812939049771907133879935,
                        104495383778385781169925874245014135708035627262536525919103632893328490496,
                        6044629098073145873860096,
                        1131377253,
                        210626190275159008588167432483938910866205453014421218013934467097,
                    ]
                ),
                28
            )
        );
}

#[test]
#[should_panic(expected: ('invalid guardian set index',))]
fn test_submit_guardian_set_rejects_wrong_index_in_signer() {
    let dispatcher = deploy(guardian_set0(), CHAIN_ID, GOVERNANCE_CHAIN_ID, GOVERNANCE_CONTRACT);

    dispatcher.submit_new_guardian_set(governance_upgrade_vm1());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm3());
}

#[test]
#[should_panic(expected: ('invalid guardian set sequence',))]
fn test_submit_guardian_set_rejects_wrong_index_in_payload() {
    let dispatcher = deploy(
        array_try_into(array![0x686b9ea8e3237110eaaba1f1b7467559a3273819]),
        CHAIN_ID,
        GOVERNANCE_CHAIN_ID,
        GOVERNANCE_CONTRACT
    );

    dispatcher
        .submit_new_guardian_set(
            ByteArrayImpl::new(
                array_try_into(
                    array![
                        1766847064779992086486657352557640859156186269544082638392748527776826036,
                        194277059768327503957595402526511955581212349039595919400671200491321209820,
                        434496814532575177274556800095069791478767471500230428914335519265903345664,
                        4835703278458516699153920,
                        1131377253,
                        210626190275159756877005745906233031152839803751327281851396470809,
                    ]
                ),
                28
            )
        );
}

#[test]
#[should_panic(expected: ('no guardians specified',))]
fn test_deploy_rejects_empty() {
    deploy(array![], CHAIN_ID, GOVERNANCE_CHAIN_ID, GOVERNANCE_CONTRACT);
}

#[test]
#[should_panic(expected: ('no guardians specified',))]
fn test_submit_guardian_set_rejects_empty() {
    let dispatcher = deploy(
        array_try_into(array![0x686b9ea8e3237110eaaba1f1b7467559a3273819]),
        CHAIN_ID,
        GOVERNANCE_CHAIN_ID,
        GOVERNANCE_CONTRACT
    );

    dispatcher
        .submit_new_guardian_set(
            ByteArrayImpl::new(
                array_try_into(
                    array![
                        1766847064779991325832211120621568385196129637393266719212748612567366167,
                        101913621687013458268034923649142123047718029322567283658602535951581008800,
                        334673952516423004763530105323893104621683276518933872430791864882642812928,
                        4835703278458516699153920,
                        1131377253,
                        144116287587483904,
                    ]
                ),
                8
            )
        );
}

#[test]
#[fuzzer(runs: 100, seed: 0)]
#[should_panic]
fn test_submit_guardian_set_rejects_corrupted(pos: usize, random1: usize, random2: usize) {
    let dispatcher = deploy(guardian_set0(), CHAIN_ID, GOVERNANCE_CHAIN_ID, GOVERNANCE_CONTRACT);

    let vm = corrupted_vm(governance_upgrade_vm1(), pos, random1, random2);
    dispatcher.submit_new_guardian_set(vm);
}

#[test]
#[should_panic(expected: ('wrong governance chain',))]
fn test_submit_guardian_set_rejects_non_governance(pos: usize, random1: usize, random2: usize) {
    let dispatcher = deploy(guardian_set0(), CHAIN_ID, GOVERNANCE_CHAIN_ID, GOVERNANCE_CONTRACT);

    dispatcher.submit_new_guardian_set(governance_upgrade_vm1());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm2());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm3());

    dispatcher.submit_new_guardian_set(good_vm1());
}

fn deploy(
    guardians: Array<EthAddress>,
    chain_id: u16,
    governance_chain_id: u16,
    governance_contract: u256,
) -> IWormholeDispatcher {
    let mut args = array![];
    (guardians, chain_id, governance_chain_id, governance_contract).serialize(ref args);
    let contract = declare("wormhole");
    let contract_address = match contract.deploy(@args) {
        Result::Ok(v) => { v },
        Result::Err(err) => {
            panic(err.panic_data);
            0.try_into().unwrap()
        },
    };
    IWormholeDispatcher { contract_address }
}

pub fn deploy_and_init() -> IWormholeDispatcher {
    let dispatcher = deploy(guardian_set0(), CHAIN_ID, GOVERNANCE_CHAIN_ID, GOVERNANCE_CONTRACT);

    dispatcher.submit_new_guardian_set(governance_upgrade_vm1());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm2());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm3());
    dispatcher.submit_new_guardian_set(governance_upgrade_vm4());

    dispatcher
}

fn corrupted_vm(mut real_data: ByteArray, pos: usize, random1: usize, random2: usize) -> ByteArray {
    let mut new_data = array![];

    // Make sure we select a position not on the last item because
    // we didn't implement corrupting an incomplete bytes31.
    let pos = pos % (real_data.len() - 31);
    let bad_index = pos / 31;

    let mut num_last_bytes = 0;
    let mut i = 0;
    loop {
        let (real_bytes, num_bytes) = match real_data.pop_front() {
            Option::Some(v) => v,
            Option::None => { break; }
        };
        if num_bytes < 31 {
            new_data.append(real_bytes);
            num_last_bytes = num_bytes;
            break;
        }

        if i == bad_index {
            new_data.append(corrupted_bytes(real_bytes, pos % 31, random1, random2));
        } else {
            new_data.append(real_bytes);
        }
        i += 1;
    };
    ByteArrayImpl::new(new_data, num_last_bytes)
}

// Returns a `bytes31` value with 2 bytes changed. We need to change at least 2 bytes
// because a single byte can be a recovery id, where only 1 bit matters so
// a modification of recovery id can result in a valid VM.
fn corrupted_bytes(input: bytes31, index: usize, random1: usize, random2: usize) -> bytes31 {
    let index2 = (index + 1) % 31;

    let mut value: u256 = 0;
    let mut i = 0;
    while i < 31 {
        let real_byte = input.at(30 - i);
        let new_byte = if i == index {
            corrupted_byte(real_byte, random1)
        } else if i == index2 {
            corrupted_byte(real_byte, random2)
        } else {
            real_byte
        };
        value = value * 256 + new_byte.into();
        i += 1;
    };
    let value: felt252 = value.try_into().expect('corrupted bytes overflow');
    value.try_into().expect('corrupted bytes overflow')
}

// Returns a byte that's not equal to `value`.
fn corrupted_byte(value: u8, random: usize) -> u8 {
    let v: u16 = value.into() + 1 + (random % 255).try_into().unwrap();
    (v % 256).try_into().unwrap()
}

fn guardian_set0() -> Array<EthAddress> {
    array_try_into(array![0x58CC3AE5C097b213cE3c81979e1B9f9570746AA5])
}

// A random VAA pulled from Hermes.
fn good_vm1() -> ByteArray {
    let bytes = array![
        1766847066033410293701000231337210964058791470455465385734308943533652138,
        250126301534699068413432844632573953364878937343368310395142095034982913232,
        374780571002258088211231890250917843593951765403462661483498298003400611238,
        23190137343211334092589308306056431640588154666326612124726174150537328574,
        238750269065878649216923353030193912502813798896051725498208457553032584635,
        29844190303057534696518006438077948796328243878877072296680853158289181326,
        106329507856770018708432343978518079724691760719405501795955774399597471533,
        50779865592261858016477142415230454208001695486195806892438697217059319645,
        448669871976126446102256476358498380455807705600424321390063431836375575318,
        115682669871397824853706713833773246708114483862317474710603223566297521279,
        301634766618012930739391408723909107532790832406455099966028276947414082504,
        104473166230846104217366042152018649207811514257244625711402436055500423094,
        64445621634231668761998815864645440965239569561546522651415024970517905416,
        192317190225976528694195501079591384434869624408066864018183189813956862386,
        289982656017597431343118552054719821766658675456661448685110903889153449006,
        218840601196095059731241556733624112758673153548932709011933806481899680620,
        430933799927481265070475198137531816946660368757134588278434352703899277070,
        69322998883710289192076494057541346430050879299268159627180404869988632073,
        23862615839737051269352321086490452186237833007444069999578906611768140646,
        444634264607471510688862284107804392707600799506487897206707262445172121289,
        438038196736233160320436150616293672539386464061037100698335568417587662951,
        4682255185797880874381673193118803274635247527626050223938224759013169366,
        337620725992972686809095065321563509600769533202700218393281926304544120094,
        106657917096532484607371891267699639824731774168349872862335217581425289654,
        71240348385993236445536577509595968468284689483611375124653855125285401592,
        347603391821038175842934311068097986460257977131947418186118379296987051086,
        414263571545410645948841360836383289766662078574048514890988877286444618669,
        250301638008739107522011802538487063969565433276260914336890309092111026583,
        43192785595291340058788190601908070333310658291317702311902081,
        52685537088250779930155363779405986390839624071318818148325576008719597568,
        14615204155786886573933667335033405822686404253588533,
    ];
    ByteArrayImpl::new(array_try_into(bytes), 22)
}

const CHAIN_ID: u16 = 1;
const GOVERNANCE_CHAIN_ID: u16 = 1;
const GOVERNANCE_CONTRACT: u256 = 4;

// Below are actual guardian set upgrade VAAS from
// https://github.com/pyth-network/pyth-crosschain/blob/main/contract_manager/src/contracts/wormhole.ts#L32-L37
fn governance_upgrade_vm1() -> ByteArray {
    let bytes = array![
        1766847064779994277746302277072294871108550301449637470263976489521154979,
        374953657095152923031303770743522269007103499920836805761143506434463979495,
        373725794026553362537846905304981854320892126869150736450761801254169477120,
        4835703278458516786446336,
        1131377253,
        3533694129556775410652111826415980944262631656421498398215501759245151417,
        145493015216602589471695207668173527044214450021182755196032581352392984224,
        267497573836069714380350521200881787609530659298168186016481773490244091266,
        443348533394886521835330696538264729103669807313401311199245411889706258110,
        200303433165499832354845293203843028338419687800279845786613090211434473108,
        37282816539161742972709083946551920068062204748477644719930149699874385462,
        111200938271608595261384934914291476246753101189480743698823215749338358345,
        5785682963869019134199015821749288033381872318410562688180948003975909269,
        372447340016996751453958019806457886428852701283870538393820846119845147788,
        33251468085387571623103303511315696691298281336333243761063342581452341650,
        323161992096034641767541451811925056802673576212351940217752194462561980347,
        55852237138651071644815135002358067220635692701051811455610533875912981641,
        190413173566657072516608762222993749133,
    ];
    ByteArrayImpl::new(array_try_into(bytes), 16)
}

fn governance_upgrade_vm2() -> ByteArray {
    let bytes = array![
        1766847065210651126944505525521222069599854288126726949998063840465138797,
        39054013088470866893599813322035661048501117089555726788687392536164861911,
        186918267413056900890218521593203545230034250983266377769400670103688217224,
        54214750922097681971590495378823998039261772575502204791108963815396679538,
        248994008232667872698758486001506749115615219491023036208140426934463230397,
        224235483823871042187452194579994340291351644933258737665365374081962645784,
        129444929990547403397151941456764812818993218972657847255833740800106200260,
        318548597134783137700627869311159664823693239605331630210202210299165477657,
        308852933010951129895011963941866000261904600807292935694851016610643657184,
        57272874228621364335170723193404742446392708606247574725663969156507973216,
        268057363923565984687253533797066429881117576606682526627284795527707196557,
        421894189151847402000239734668088381785344768284464330085711322870200424121,
        387584417395337067059819722404321580961380603778956902593594676080614899975,
        443523131755342451570503958659975367050637907361274652611595274647186073286,
        375107813087591446268414166006799680982485924290770541964399283524664437852,
        269085314426821465871247165234113878276592898426373369094610290261431112145,
        394348694527460459816454348661249546781513091938003106009521096332701847735,
        125670844183465056159554034633959680574572737212268749779705635378223489518,
        35053869475614771227400345921974210525173525784259894123687028214330135561,
        57544237843860512274491447149877290860624174166427313971286819807773907946,
        110681468147560408039747352809294082929396639199393789980441736520638055418,
        45709869872872997180086402576677640909777820941436708911954532772405320395,
        339910511168418517917975736269171135752028257685502872671902330279073260362,
        76482764517489607955778008190826845581092838692650194719207882266659669490,
        443663869577220861680293443959666949893574779475752540587040489501289361777,
        231057874919577223790659840464529054850239319545221055959374917590157019925,
        309066533217885002074480704480667619809952056265738927105682076502747220549,
        212379788814604791028007106821871908074818251907335322546331543385945316762,
        165408661499085325620077702639227003047567884011538988728381864749733773312,
        29852013947978990147012099107546124222651092940097518043136,
        5874446556610227706402640854088357165514903,
        314635470832203526600706472223155382046271943513513368538979543914002758100,
        289993023590817330918274026889451152915026890048318491140264484864242055689,
        211265316833000774821515110003986084297271807500310630074520699505436206838,
        314620948986744608212517578488307826224331071350776523303159889004405167502,
        242768143829057016675085776170635413106817756232919004913342240722183648628,
        289318220340670045883106021427202666948428587921558828582664470923483208386,
        254304247593881109676724582609273741670949040469906895867342151706444640548,
        324707984060675446628128892371664948354047882542253609514703956739624414429,
        125786084546320950738753348592393927755418642173185609412108154831520915923,
        192033422676298173731756291271054199566981168481817292625435767748408605264,
        70237018464728620254434305961956673950089621204502627373468857093940647376,
        75218391584551901010047495874303520775865073092730040058902770251005073864,
        13453,
    ];
    ByteArrayImpl::new(array_try_into(bytes), 2)
}

fn governance_upgrade_vm3() -> ByteArray {
    let bytes = array![
        1766847065622031860560134801367788401015571316785630090767859240961980367,
        408239335069601434456324970231880063141100099721451058487412176729277688481,
        239499689753305520381601038928513796227317320911002802551444547638809838552,
        377515301744513788082523380503265415136454305699441419871183462725292421897,
        293792427782284265052215516935690252661713253004854348809536189867737815900,
        307592266914902727916633292174670243188255704510048236277225897328900269063,
        127373290139474278928992577974696343462455769904046193125018730285162391253,
        391788800785481654990215164673817619378887263909639120513493454202816019863,
        410413307118599096084169722886408584518140871169074821252461819158667354254,
        18837648490111595970198935552266546643395427923804444528866768515056419823,
        29964034682984173558839379357821569529808274426015494950430380078317881569,
        86017868501532670528023530422115758730056738654625156800662778409971102579,
        316587967137295297243489759859408971025377360462781809717927347025414193161,
        412080542369789462767669836400697110505430973769191785499739175360336337147,
        342817079347905714229318925597762381802367663565411998187223317628701911440,
        323381353160339090772037140072061985169258958022395380273676898316834639836,
        199931590715070935127318740956564588449721873695471932311700469202637695100,
        53310522180389647586576928116330850824055549848985438538201222342553700451,
        387322343922164253479438966163491855981414317104760621828688810466847848718,
        81609701542274539489711635515209037026645626576756528749469616228397567798,
        182108205861564989333892774796475580877981373947799860454217397980367659628,
        21549663410658134468902761710868642321546772465973442277960059676129502668,
        189434039785735939400321781125039794740638779195156759980704929066694157130,
        52255833533187953003213955242027099690232530588872309460610106220279805641,
        197105018621162723372171195601447549272902142615124680111298974553437412361,
        243585516151555343004264928593678764289083751554802049062044286334698216184,
        98577806073901898829375415748245478967425496216912736575886605480181121443,
        92916551389967933235240931764170084503123511470557201449603712010638670912,
        279247190794260779926452059689914005511524938154821508635388069101252378624,
        27765181507524306000048567556593270127801507143251178553344,
        5874446556610227706402640926145951203442839,
        314635470832203526600706472223155382046271943513513368538979543914002758100,
        289993023590817330918274026889451152915026890048318491140264484864242055689,
        211265316833000774821515110003986084297271807500310630074520699505436206838,
        314620948986744608212517578488307826224331071350776523303159889004405167502,
        242768143829057016675085658054156069029173843566452718977789980910319968372,
        289318220340670045883106021427202666948428587921558828582664470923483208386,
        254304247593881109676724582609273741670949040469906895867342151706444640548,
        324707984060675446628128892371664948354047882542253609514703956739624414429,
        125786084546320950738753348592393927755418642173185609412108154831520915923,
        192033422676298173731756291271054199566981168481817292625435767748408605264,
        70237018464728620254434305961956673950089621204502627373468857093940647376,
        75218391584551901010047495874303520775865073092730040058902770251005073864,
        13453,
    ];
    ByteArrayImpl::new(array_try_into(bytes), 2)
}

fn governance_upgrade_vm4() -> ByteArray {
    let bytes = array![
        1766847066033426987337757245669159273063358729535478806850006662056807068,
        191023158244075433218055661747029015323596061316379687901032561397223546211,
        30156550775609732785124128084945604136341937408029517653427049258063209215,
        301841618969457377999846355946508544313297803407604349411042057045510372286,
        399879387152070823070522891203531321261797829310211644637928969034931151834,
        1184971666775858810527395126763859219514013163556756790208661779020321698,
        427827873217506136303198988655697899138087317492051993053159867826892618987,
        55439109913191967501571602277685262841453050617358377329061538069328212552,
        34944602254693785869427132065664922748183924456022812505745784482260734500,
        50091615215549712387991200985741575718080363004681463525186508796585379155,
        265247833149227842278059484961926330281584344437952973839486092319885300192,
        421631446041795295328070636491346018953171276542115189082171457479754499396,
        59057903625576869235407103565877017330396402246452653660114888284545941770,
        315797852826246435174946736461540321579373154656484006452063031513301027405,
        9521420622979958910372839981791309896262328383324674284772682980734269170,
        272964069264268937653695089515793248726920319914036642027008415285406913245,
        194708434228888226032102758315234166672190899487218971410889219134701358728,
        117864954129109327302856065706421701676973859697066630532570005860486924893,
        323457021720552374478769194145226061243431674370101604382965685057422991463,
        327482733702858147057975319784026874245182397914737119038454598086198587150,
        159726033816658034104416471293601013976445904149240898589368461412472508473,
        165970343982649234398221341351816767302457220504375238905210573566962780340,
        66269488760319836583658182431744051236825244016843316092957806563966254500,
        360882001282595740056823749884962228392082962172369522212117195988772429063,
        202692667772209236945884489592750537635169234501360011152939202347962132650,
        407257364829649465305126488148712878739144584682351279109461295389594525334,
        270499607712829989691415988895838806019492861138165540862008308077962735002,
        388443296961168536186587069708212659389994895697827691755155284015603161464,
        45068266527940236008536134081672474027695203549460934893262212861351952384,
        31319268777966350508118557206583844424308993254125039779840,
        5874446556610227706402640998203302487747647,
        204224545225244051821590480758420624947979343122083461045877549162059250132,
        289993023590817330918274026889451152915026890048318491140264484864242055689,
        211265316833000774821515110003986084297271807500310630074520699505436206838,
        314620948986744608212517578488307826224331071350776523303159889004405167502,
        242768143829057016675085658054156069029173843566452718977789980910319968372,
        289318220340670045883106021427202666948428587921558828582664470923483208386,
        254304247593881109676724582609273741670949040469906895867342151706444640548,
        324707984060675446628128892371664948354047882542253609514703956739624414429,
        125786084546320950738753348592393927755418642173185609412108154831520915923,
        192033422676298173731756291271054199566981168481817292625435767748408605264,
        70237018464728620254434305961956673950089621204502627373468857093940647376,
        75218391584551901010047495874303520775865073092730040058902770251005073864,
        13453,
    ];
    ByteArrayImpl::new(array_try_into(bytes), 2)
}