#include "imports/stdlib.fc";
#include "common/errors.fc";
#include "common/storage.fc";
#include "common/utils.fc";
#include "common/constants.fc";
#include "common/merkle_tree.fc";
#include "common/governance_actions.fc";
#include "common/gas.fc";
#include "common/op.fc";
#include "./Wormhole.fc";

cell store_price(int price, int conf, int expo, int publish_time) {
    return begin_cell()
        .store_int(price, 64)
        .store_uint(conf, 64)
        .store_int(expo, 32)
        .store_uint(publish_time, 64)
    .end_cell();
}

slice read_and_verify_header(slice data) impure {
    int magic = data~load_uint(32);
    throw_unless(ERROR_INVALID_MAGIC, magic == ACCUMULATOR_MAGIC);
    int major_version = data~load_uint(8);
    throw_unless(ERROR_INVALID_MAJOR_VERSION, major_version == MAJOR_VERSION);
    int minor_version = data~load_uint(8);
    throw_if(ERROR_INVALID_MINOR_VERSION, minor_version < MINIMUM_ALLOWED_MINOR_VERSION);
    int trailing_header_size = data~load_uint(8);
    ;; skip trailing headers
    data~skip_bits(trailing_header_size * 8);
    int update_type = data~load_uint(8);
    throw_unless(ERROR_INVALID_UPDATE_DATA_TYPE, update_type == WORMHOLE_MERKLE_UPDATE_TYPE);
    return data;
}

(int, int, int, int, int, int, int, int, slice) read_and_verify_message(slice cs, int root_digest) impure {
    int message_size = cs~load_uint(16);
    (cell message, slice cs) = read_and_store_large_data(cs, message_size * 8);
    slice message = message.begin_parse();
    slice cs = read_and_verify_proof(root_digest, message, cs);

    int message_type = message~load_uint(8);
    throw_unless(ERROR_INVALID_MESSAGE_TYPE, message_type == PRICE_FEED_MESSAGE_TYPE);

    int price_id = message~load_uint(256);
    int price = message~load_int(64);
    int conf = message~load_uint(64);
    int expo = message~load_int(32);
    int publish_time = message~load_uint(64);
    int prev_publish_time = message~load_uint(64);
    int ema_price = message~load_int(64);
    int ema_conf = message~load_uint(64);

    return (price_id, price, conf, expo, publish_time, prev_publish_time, ema_price, ema_conf, cs);
}

(int, int, int, int) parse_price(slice price_feed) {
    int price = price_feed~load_int(64);
    int conf = price_feed~load_uint(64);
    int expo = price_feed~load_int(32);
    int publish_time = price_feed~load_uint(64);
    return (price, conf, expo, publish_time);
}

int get_update_fee(slice data) method_id {
    load_data();
    slice cs = read_and_verify_header(data);
    int wormhole_proof_size_bytes = cs~load_uint(16);
    (cell wormhole_proof, slice cs) = read_and_store_large_data(cs, wormhole_proof_size_bytes * 8);
    int num_updates = cs~load_uint(8);
    return single_update_fee * num_updates;
}

int get_single_update_fee() method_id {
    load_data();
    return single_update_fee;
}

int get_governance_data_source_index() method_id {
    load_data();
    return governance_data_source_index;
}

cell get_governance_data_source() method_id {
    load_data();
    return governance_data_source;
}

int get_last_executed_governance_sequence() method_id {
    load_data();
    return last_executed_governance_sequence;
}

int get_is_valid_data_source(cell data_source) method_id {
    load_data();
    int data_source_key = cell_hash(data_source);
    (slice value, int found?) = is_valid_data_source.udict_get?(256, data_source_key);
    if (found?) {
        return value~load_int(1);
    } else {
        return 0;
    }
}

(int, int, int, int) get_price_unsafe(int price_feed_id) method_id {
    load_data();
    (slice result, int success) = latest_price_feeds.udict_get?(256, price_feed_id);
    throw_unless(ERROR_PRICE_FEED_NOT_FOUND, success);
    slice price_feed = result~load_ref().begin_parse();
    slice price = price_feed~load_ref().begin_parse();
    return parse_price(price);
}

(int, int, int, int) get_price_no_older_than(int time_period, int price_feed_id) method_id {
    load_data();
    (int price, int conf, int expo, int publish_time) = get_price_unsafe(price_feed_id);
    int current_time = now();
    throw_if(ERROR_OUTDATED_PRICE, max(0, current_time - publish_time) > time_period);
    return (price, conf, expo, publish_time);
}

(int, int, int, int) get_ema_price_unsafe(int price_feed_id) method_id {
    load_data();
    (slice result, int success) = latest_price_feeds.udict_get?(256, price_feed_id);
    throw_unless(ERROR_PRICE_FEED_NOT_FOUND, success);
    slice price_feed = result~load_ref().begin_parse();
    slice price = price_feed~load_ref().begin_parse();
    slice ema_price = price_feed~load_ref().begin_parse();
    return parse_price(ema_price);
}

(int, int, int, int) get_ema_price_no_older_than(int time_period, int price_feed_id) method_id {
    load_data();
    (int price, int conf, int expo, int publish_time) = get_ema_price_unsafe(price_feed_id);
    int current_time = now();
    throw_if(ERROR_OUTDATED_PRICE, max(0, current_time - publish_time) > time_period);
    return (price, conf, expo, publish_time);
}

(int, int) parse_data_source(cell data_source) {
    slice ds = data_source.begin_parse();
    int emitter_chain = ds~load_uint(16);
    int emitter_address = ds~load_uint(256);
    return (emitter_chain, emitter_address);
}

int parse_pyth_payload_in_wormhole_vm(slice payload) impure {
    int accumulator_wormhole_magic = payload~load_uint(32);
    throw_unless(ERROR_INVALID_MAGIC, accumulator_wormhole_magic == ACCUMULATOR_WORMHOLE_MAGIC);

    int update_type = payload~load_uint(8);
    throw_unless(ERROR_INVALID_UPDATE_DATA_TYPE, update_type == WORMHOLE_MERKLE_UPDATE_TYPE);

    payload~load_uint(64);  ;; Skip slot
    payload~load_uint(32);  ;; Skip ring_size

    return payload~load_uint(160);  ;; Return root_digest
}

() calculate_and_validate_fees(int msg_value, int num_updates) impure {
    int update_fee = single_update_fee * num_updates;
    int compute_fee = get_compute_fee(
        WORKCHAIN,
        UPDATE_PRICE_FEEDS_BASE_GAS + (UPDATE_PRICE_FEEDS_PER_UPDATE_GAS * num_updates)
    );
    throw_unless(ERROR_INSUFFICIENT_GAS, msg_value >= compute_fee);
    int remaining_msg_value = msg_value - compute_fee;

    ;; Check if the sender has sent enough TON to cover the update_fee
    throw_unless(ERROR_INSUFFICIENT_FEE, remaining_msg_value >= update_fee);
}

(int) find_price_id_index(tuple price_ids, int price_id) {
    int len = price_ids.tlen();
    int i = 0;
    while (i < len) {
        if (price_ids.at(i) == price_id) {
            return i;
        }
        i += 1;
    }
    return -1;  ;; Not found
}


tuple parse_price_feeds_from_data(int msg_value, slice data, tuple price_ids, int min_publish_time, int max_publish_time, int unique) {
    slice cs = read_and_verify_header(data);

    int wormhole_proof_size_bytes = cs~load_uint(16);
    (cell wormhole_proof, slice new_cs) = read_and_store_large_data(cs, wormhole_proof_size_bytes * 8);
    cs = new_cs;

    int num_updates = cs~load_uint(8);

    calculate_and_validate_fees(msg_value, num_updates);

    (_, _, _, _, int emitter_chain_id, int emitter_address, _, _, slice payload, _) = parse_and_verify_wormhole_vm(wormhole_proof.begin_parse());

    ;; Check if the data source is valid
    cell data_source = begin_cell()
        .store_uint(emitter_chain_id, 16)
        .store_uint(emitter_address, 256)
    .end_cell();

    ;; Dictionary doesn't support cell as key, so we use cell_hash to create a 256-bit integer key
    int data_source_key = cell_hash(data_source);
    (slice value, int found?) = is_valid_data_source.udict_get?(256, data_source_key);
    throw_unless(ERROR_UPDATE_DATA_SOURCE_NOT_FOUND, found?);
    int valid = value~load_int(1);
    throw_unless(ERROR_INVALID_UPDATE_DATA_SOURCE, valid);

    int root_digest = parse_pyth_payload_in_wormhole_vm(payload);

    ;; Create dictionary to store price feeds in order (dict has a udict_get_next? method which returns the next key in order)
    cell ordered_feeds = new_dict();
    ;; Track which price IDs we've found
    cell found_price_ids = new_dict();

    int index = 0;

    repeat(num_updates) {
        (int price_id, int price, int conf, int expo, int publish_time, int prev_publish_time, int ema_price, int ema_conf, slice new_cs) = read_and_verify_message(cs, root_digest);
        cs = new_cs;

        int price_ids_len = price_ids.tlen();

        ;; Check if we've already processed this price_id to avoid duplicates
        (_, int already_processed?) = found_price_ids.udict_get?(256, price_id);
        if (~ already_processed?) { ;; Only process if we haven't seen this price_id yet
            int should_include = (price_ids_len == 0)
                | ((price_ids_len > 0)
                & (publish_time >= min_publish_time)
                & (publish_time <= max_publish_time)
                & ((unique == 0) | (min_publish_time > prev_publish_time)));

            if (should_include) {
                ;; Create price feed cell containing both current and EMA prices
                cell price_feed_cell = begin_cell()
                    .store_ref(store_price(price, conf, expo, publish_time))
                    .store_ref(store_price(ema_price, ema_conf, expo, publish_time))
                .end_cell();

                if (price_ids_len == 0) {
                    ordered_feeds~udict_set(8, index, begin_cell()
                        .store_uint(price_id, 256)
                        .store_ref(price_feed_cell)
                    .end_cell().begin_parse());
                    index += 1;
                } else {
                    index = find_price_id_index(price_ids, price_id);
                    if (index >= 0) {
                        ordered_feeds~udict_set(8, index, begin_cell()
                            .store_uint(price_id, 256)
                            .store_ref(price_feed_cell)
                        .end_cell().begin_parse());
                    }
                }

                ;; Mark this price ID as found
                found_price_ids~udict_set(256, price_id, begin_cell().store_int(true, 1).end_cell().begin_parse());
            }
        }
    }

    throw_if(ERROR_INVALID_UPDATE_DATA_LENGTH, ~ cs.slice_empty?());

    ;; Verify all requested price IDs were found
    if (price_ids.tlen() > 0) {
        int i = 0;
        repeat(price_ids.tlen()) {
            int requested_id = price_ids.at(i);
            (_, int found?) = found_price_ids.udict_get?(256, requested_id);
            throw_unless(ERROR_PRICE_FEED_NOT_FOUND_WITHIN_RANGE, found?);
            i += 1;
        }
    }

    ;; Create final ordered tuple from dictionary
    tuple price_feeds = empty_tuple();
    int index = -1;
    do {
        (index, slice value, int success) = ordered_feeds.udict_get_next?(8, index);
        if (success) {
            tuple price_feed = empty_tuple();
            price_feed~tpush(value~load_uint(256));  ;; price_id
            price_feed~tpush(value~load_ref());      ;; price_feed_cell
            price_feeds~tpush(price_feed);
        }
    } until (~ success);

    return price_feeds;
}

;; Creates a chain of cells from price feeds, with each cell containing exactly one price_id (256 bits)
;; and one ref to the price feed cell. Returns the head of the chain.
;; Each cell now contains exactly:
;; - One price_id (256 bits)
;; - One ref to price_feed_cell
;; - One optional ref to next cell in chain
;; This approach is:
;; - More consistent with TON's cell model
;; - Easier to traverse and read individual price feeds
;; - Cleaner separation of data
;; - More predictable in terms of cell structure
cell create_price_feed_cell_chain(tuple price_feeds) {
    cell result = null();

    int i = price_feeds.tlen() - 1;
    while (i >= 0) {
        tuple price_feed = price_feeds.at(i);
        int price_id = price_feed.at(0);
        cell price_feed_cell = price_feed.at(1);

        ;; Create new cell with single price feed and chain to previous result
        builder current_builder = begin_cell()
            .store_uint(price_id, 256)    ;; Store price_id
            .store_ref(price_feed_cell);   ;; Store price data ref

        ;; Chain to previous cells if they exist
        if (~ cell_null?(result)) {
            current_builder = current_builder.store_ref(result);
        }

        result = current_builder.end_cell();
        i -= 1;
    }

    return result;
}

() send_price_feeds_response(tuple price_feeds, int msg_value, int op, slice sender_address, slice target_address, slice custom_payload) impure {
    ;; Build response cell with price feeds
    builder response = begin_cell()
        .store_uint(op, 32)  ;; Response op
        .store_uint(price_feeds.tlen(), 8);  ;; Number of price feeds

    ;; Create and store price feed cell chain
    cell price_feeds_cell = create_price_feed_cell_chain(price_feeds);
    cell custom_payload_cell = begin_cell().store_slice(custom_payload).end_cell();
    response = response.store_ref(price_feeds_cell).store_slice(sender_address).store_ref(custom_payload_cell);

    int num_price_feeds = price_feeds.tlen();

    ;; Calculate all fees
    int compute_fee = get_compute_fee(WORKCHAIN, get_gas_consumed());
    int update_fee = single_update_fee * price_feeds.tlen();

    ;; Calculate total fees and remaining excess
    int total_fees = compute_fee + update_fee;
    int excess = msg_value - total_fees;

    ;; SECURITY: Integrators MUST validate that messages are from this Pyth contract
    ;; in their receive function. Otherwise, attackers could:
    ;; 1. Send invalid price responses
    ;; 2. Impersonate users via sender_address and custom_payload fields
    ;; 3. Potentially drain the protocol
    ;;
    ;; Note: This message is bounceable. If the target contract rejects the message,
    ;; the excess TON will remain in this contract and won't be automatically refunded to the
    ;; original sender. Integrators should handle failed requests and refunds in their implementation.
    send_raw_message(begin_cell()
        .store_uint(0x18, 6)
        .store_slice(target_address)
        .store_coins(excess)
        .store_uint(1, MSG_SERIALIZE_BITS)
        .store_ref(response.end_cell())
        .end_cell(),
        0);
}

() parse_price_feed_updates(int msg_value, slice update_data_slice, slice price_ids_slice, int min_publish_time, int max_publish_time, slice sender_address, slice target_address, slice custom_payload) impure {
    load_data();

    ;; Load price_ids tuple
    int price_ids_len = price_ids_slice~load_uint(8);
    tuple price_ids = empty_tuple();
    repeat(price_ids_len) {
        int price_id = price_ids_slice~load_uint(256);
        price_ids~tpush(price_id);
    }

    tuple price_feeds = parse_price_feeds_from_data(msg_value, update_data_slice, price_ids, min_publish_time, max_publish_time, false);
    send_price_feeds_response(price_feeds, msg_value, OP_PARSE_PRICE_FEED_UPDATES, sender_address, target_address, custom_payload);
}

() parse_unique_price_feed_updates(int msg_value, slice update_data_slice, slice price_ids_slice, int publish_time, int max_staleness, slice sender_address, slice target_address, slice custom_payload) impure {
    load_data();

    ;; Load price_ids tuple
    int price_ids_len = price_ids_slice~load_uint(8);
    tuple price_ids = empty_tuple();
    repeat(price_ids_len) {
        int price_id = price_ids_slice~load_uint(256);
        price_ids~tpush(price_id);
    }

    tuple price_feeds = parse_price_feeds_from_data(msg_value, update_data_slice, price_ids, publish_time, publish_time + max_staleness, true);
    send_price_feeds_response(price_feeds, msg_value, OP_PARSE_UNIQUE_PRICE_FEED_UPDATES, sender_address, target_address, custom_payload);
}

() update_price_feeds(int msg_value, slice data) impure {
    load_data();
    tuple price_feeds = parse_price_feeds_from_data(msg_value, data, empty_tuple(), 0, 0, false);
    int num_updates = price_feeds.tlen();

    int i = 0;
    while(i < num_updates) {
        tuple price_feed = price_feeds.at(i);
        int price_id = price_feed.at(0);
        cell price_feed_cell = price_feed.at(1);
        slice price_feed = price_feed_cell.begin_parse();
        slice price = price_feed~load_ref().begin_parse();
        slice ema_price = price_feed~load_ref().begin_parse();
        (int price_, int conf, int expo, int publish_time) = parse_price(price);

        (slice latest_price_info, int found?) = latest_price_feeds.udict_get?(256, price_id);
        int latest_publish_time = 0;
        if (found?) {
            slice price_feed_slice = latest_price_info~load_ref().begin_parse();
            slice price_slice = price_feed_slice~load_ref().begin_parse();

            price_slice~load_int(64);  ;; Skip price
            price_slice~load_uint(64); ;; Skip conf
            price_slice~load_int(32);  ;; Skip expo
            latest_publish_time = price_slice~load_uint(64);
        }

        if (publish_time > latest_publish_time) {
            latest_price_feeds~udict_set(256, price_id, begin_cell().store_ref(price_feed_cell).end_cell().begin_parse());
        }
        i += 1;
    }

    store_data();
}

() verify_governance_vm(int emitter_chain_id, int emitter_address, int sequence) impure {
    (int gov_chain_id, int gov_address) = parse_data_source(governance_data_source);
    throw_unless(ERROR_INVALID_GOVERNANCE_DATA_SOURCE, emitter_chain_id == gov_chain_id);
    throw_unless(ERROR_INVALID_GOVERNANCE_DATA_SOURCE, emitter_address == gov_address);
    throw_if(ERROR_OLD_GOVERNANCE_MESSAGE, sequence <= last_executed_governance_sequence);
    last_executed_governance_sequence = sequence;
}

(int, int, slice) parse_governance_instruction(slice payload) impure {
    int magic = payload~load_uint(32);
    throw_unless(ERROR_INVALID_GOVERNANCE_MAGIC, magic == GOVERNANCE_MAGIC);

    int module = payload~load_uint(8);
    throw_unless(ERROR_INVALID_GOVERNANCE_MODULE, module == GOVERNANCE_MODULE);

    int action = payload~load_uint(8);

    int target_chain_id = payload~load_uint(16);

    return (target_chain_id, action, payload);
}

int apply_decimal_expo(int value, int expo) {
    int result = value;
    repeat (expo) {
        result *= 10;
    }
    return result;
}

() execute_upgrade_contract(cell new_code) impure {
    load_data();
    int hash_code = cell_hash(new_code);
    throw_unless(ERROR_INVALID_CODE_HASH, upgrade_code_hash == hash_code);

    ;; Set the new code
    set_code(new_code);

    ;; Set the code continuation to the new code
    set_c3(new_code.begin_parse().bless());

    ;; Reset the upgrade code hash
    upgrade_code_hash = 0;

    ;; Store the data to persist the reset above
    store_data();

    ;; Throw an exception to end the current execution
    ;; The contract will be restarted with the new code
    throw(0);
}

() execute_authorize_upgrade_contract(slice payload) impure {
    int code_hash = payload~load_uint(256);
    upgrade_code_hash = code_hash;
}

() execute_authorize_governance_data_source_transfer(slice payload) impure {
    ;; Verify the claim VAA
    (_, _, _, _, int emitter_chain_id, int emitter_address, int sequence, _, slice claim_payload, _) = parse_and_verify_wormhole_vm(payload);

    ;; Parse the claim payload
    (int target_chain_id, int action, slice claim_payload) = parse_governance_instruction(claim_payload);

    ;; Verify that this is a valid governance action for this chain
    throw_if(ERROR_INVALID_GOVERNANCE_TARGET, (target_chain_id != 0) & (target_chain_id != chain_id));
    throw_unless(ERROR_INVALID_GOVERNANCE_ACTION, action == REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER);

    ;; Extract the new governance data source index from the claim payload
    int new_governance_data_source_index = claim_payload~load_uint(32);

    ;; Verify that the new index is greater than the current index
    int current_index = governance_data_source_index;
    throw_if(ERROR_OLD_GOVERNANCE_MESSAGE, new_governance_data_source_index <= current_index);

    ;; Update the governance data source
    governance_data_source = begin_cell()
        .store_uint(emitter_chain_id, 16)
        .store_uint(emitter_address, 256)
    .end_cell();

    governance_data_source_index = new_governance_data_source_index;

    ;; Update the last executed governance sequence
    last_executed_governance_sequence = sequence;
}

() execute_set_data_sources(slice payload) impure {
    int num_sources = payload~load_uint(8);
    cell new_data_sources = new_dict();

    repeat(num_sources) {
        (cell data_source, slice new_payload) = read_and_store_large_data(payload, 272); ;; 272 = 256 + 16
        payload = new_payload;
        slice data_source_slice = data_source.begin_parse();
        int emitter_chain_id = data_source_slice~load_uint(16);
        int emitter_address = data_source_slice~load_uint(256);
        cell data_source = begin_cell()
            .store_uint(emitter_chain_id, 16)
            .store_uint(emitter_address, 256)
        .end_cell();
        int data_source_key = cell_hash(data_source);
        new_data_sources~udict_set(256, data_source_key, begin_cell().store_int(true, 1).end_cell().begin_parse());
    }

    ;; Verify that all data in the payload was processed
    throw_unless(ERROR_INVALID_PAYLOAD_LENGTH, payload.slice_empty?());

    is_valid_data_source = new_data_sources;
}

() execute_set_fee(slice payload) impure {
    int value = payload~load_uint(64);
    int expo = payload~load_uint(64);
    int new_fee = apply_decimal_expo(value, expo);
    single_update_fee = new_fee;
}

() execute_governance_payload(int action, slice payload) impure {
    if (action == AUTHORIZE_UPGRADE_CONTRACT) {
        execute_authorize_upgrade_contract(payload);
    } elseif (action == AUTHORIZE_GOVERNANCE_DATA_SOURCE_TRANSFER) {
        execute_authorize_governance_data_source_transfer(payload);
    } elseif (action == SET_DATA_SOURCES) {
        execute_set_data_sources(payload);
    } elseif (action == SET_FEE) {
        execute_set_fee(payload);
    } elseif (action == SET_VALID_PERIOD) {
        ;; Unsupported governance action
        throw(ERROR_INVALID_GOVERNANCE_ACTION);
    } elseif (action == REQUEST_GOVERNANCE_DATA_SOURCE_TRANSFER) {
        ;; RequestGovernanceDataSourceTransfer can only be part of
        ;; AuthorizeGovernanceDataSourceTransfer message
        throw(ERROR_INVALID_GOVERNANCE_ACTION);
    } else {
        throw(ERROR_INVALID_GOVERNANCE_ACTION);
    }
}

() execute_governance_action(slice in_msg_body) impure {
    load_data();

    (_, _, _, _, int emitter_chain_id, int emitter_address, int sequence, _, slice payload, _) = parse_and_verify_wormhole_vm(in_msg_body);

    verify_governance_vm(emitter_chain_id, emitter_address, sequence);

    (int target_chain_id, int action, slice payload) = parse_governance_instruction(payload);

    throw_if(ERROR_INVALID_GOVERNANCE_TARGET, (target_chain_id != 0) & (target_chain_id != chain_id));

    execute_governance_payload(action, payload);

    store_data();
}