package com.alterdekim.flash.decompiler;

import com.alterdekim.flash.decompiler.action.*;
import com.alterdekim.flash.decompiler.item.*;
import com.alterdekim.flash.decompiler.tag.DoAction;
import com.alterdekim.flash.decompiler.tag.ShockwaveTag;
import com.alterdekim.flash.decompiler.tag.TagType;
import com.alterdekim.flash.decompiler.util.ActionField;
import com.alterdekim.flash.decompiler.util.Bits;
import com.alterdekim.flash.decompiler.util.PayloadCompression;
import com.alterdekim.flash.decompiler.util.StringUtils;
import lombok.extern.slf4j.Slf4j;
import org.javatuples.Triplet;

import java.io.*;
import java.nio.file.Files;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.BitSet;
import java.util.List;

import static com.alterdekim.flash.decompiler.util.ByteUtil.bytesToInt;
import static com.alterdekim.flash.decompiler.util.ByteUtil.bytesToIntRev;
import static com.alterdekim.flash.decompiler.util.Compression.decompressDEFLATE;
import static com.alterdekim.flash.decompiler.util.Compression.decompressLZMA;

@Slf4j
public class FlashDecompiler {

    private PayloadCompression compression;

    public ShockwaveFile loadFromFile( File flashFile) throws IOException {
        return processSWF(Files.readAllBytes(flashFile.toPath()));
    }

    private ShockwaveFile processSWF(byte[] swfBytes) throws IOException {
        compression = PayloadCompression.getByTag(swfBytes[0]);
        byte[] data = Arrays.copyOfRange(swfBytes, 8, swfBytes.length);
        switch (compression) {
            case DEFLATE:
                data = decompressDEFLATE(data);
                break;
            case LZMA:
                data = decompressLZMA(data);
                break;
        }
        return processPayload(data);
    }

    private ShockwaveFile processPayload(byte[] payload) {
        BitSet bitSet = BitSet.valueOf(payload);
        int c = (payload.length * 8);
        int i = 0;
        int cnt = Bits.convert(bitSet.get(i, i+7), 7, false); // Rect x=[5] [x] [x] [x] [x]
        i += 7;
        i += cnt * 4;
        i = Bits.ceilToByte(i);
        float frameRate = Bits.convertP8(bitSet.get(i, i+16));
        i += 16;
        int frameCount = Bits.convert(bitSet.get(i, i+16), 16, true);
        i += 16;
        // 16-bit + 16-bit = 32 bits for each header tag.
        i = Bits.ceilToByte(i);
        List<ShockwaveTag> shockwaveTags = new ArrayList<>();
        while( i < c ) {
            BitSet tag = bitSet.get(i, i+16);
            int v = bytesToInt(new byte[] {(byte) Bits.convert(tag.get(8, 16), 8, true),
                    (byte) Bits.convert(tag.get(0, 8), 8, true)});
            i += 16;
            int type = (v & 65472) >> 6; // type
            int len = v & 63; // length in bytes
            if( len == 63 ) {
                len = Bits.convert(bitSet.get(i, i+32), 32, true);
                i += 32;
            }
            TagType tagType = TagType.fromId(type);
            if( tagType == TagType.DoAction ) {
                int g = i / 8;
                shockwaveTags.add(processDoActionTag(Arrays.copyOfRange(payload, g, g+len)));
            } else {
                log.error("Unknown tag type: {}", type);
            }
            i += len * 8;
        }
        return new ShockwaveFile(frameRate, frameCount, shockwaveTags, compression);
    }

    private DoAction processDoActionTag(byte[] data) {
        int i = 0;
        DoAction tag = new DoAction();
        tag.setActions(new ArrayList<>());
        while( i < data.length ) {
            switch(ActionField.fromId(data[i])) {
                case ConstantPool:
                    i++;
                    i = parseActionPool(data, i, tag);
                    break;
                case Push:
                    i = parseActionPush(data, i, tag);
                    break;
                case InitObject:
                    tag.getActions().add(new ActionInitObject());
                    i++;
                    break;
                case DefineLocal:
                    i++;
                    tag.getActions().add(new ActionDefineLocal());
                    break;
                case SetMember:
                    i++;
                    tag.getActions().add(new ActionSetMember());
                    break;
                case GetVariable:
                    i++;
                    tag.getActions().add(new ActionGetVariable());
                    break;
                default:
                    log.error("Unknown action field type: {}", StringUtils.bytesToHex(new byte[] {data[i]}));
                    i++;
                    break;
            }
        }
        return tag;
    }

    private Triplet<byte[], Integer, Integer> findNull(byte[] data, int i, int u) {
        int s = i;
        while( data[i] != 0x00 ) {
            i++;
            u++;
        }
        return Triplet.with(Arrays.copyOfRange(data, s, i), i, u);
    }

    private int parseActionPush(byte[] data, int i, DoAction tag) {
        i++;
        int len_in_bytes = bytesToIntRev(Arrays.copyOfRange(data, i, i+2));
        int u = 0;
        i += 2;
        int index;
        ActionPush push = new ActionPush();
        push.setItems(new ArrayList<>());
        while( u < len_in_bytes ) {
            ActionPushType type = ActionPushType.fromType(data[i]);
            switch(type) {
                case STRING:
                    i++;
                    var result = findNull(data, i, u);
                    push.getItems().add(new StringItem(new String(result.getValue0())));
                    i = result.getValue1();
                    u = result.getValue2();
                    break;
                case POOL_INDEX_BIG:
                    i++;
                    index = bytesToIntRev(Arrays.copyOfRange(data, i, i+2));
                    push.getItems().add(new PoolIndexItem(index));
                    i += 2;
                    u += 3;
                    continue;
                case POOL_INDEX:
                    i++;
                    index = bytesToIntRev(new byte[] {data[i]});
                    push.getItems().add(new PoolIndexItem(index));
                    i++;
                    u+=2;
                    continue;
                case INTEGER:
                    i++;
                    push.getItems().add(new IntegerItem(bytesToIntRev(Arrays.copyOfRange(data, i, i+4))));
                    i += 4;
                    u += 5;
                    continue;
                case NULL:
                    push.getItems().add(new NullItem());
                    i++;
                    u++;
                    continue;
                default:
                    i++;
                    u++;
                    log.info("parseActionPush(): Unknown action field: {}", data[i]);
            }
        }
        tag.getActions().add(push);
        return i;
    }

    private int parseActionPool(byte[] data, int i, DoAction tag) {
        List<String> actionPool = new ArrayList<>();
        int bytes_len = bytesToIntRev(Arrays.copyOfRange(data, i, i+2));
        i += 2;
        int count = bytesToIntRev(Arrays.copyOfRange(data, i, i+2));
        i += 2;
        int u = 1;
        StringBuilder s = new StringBuilder();
        while( u <= count ) {
            if( data[i] == 0x00 ) {
                actionPool.add(s.toString());
                s = new StringBuilder();
                u++;
                i++;
                continue;
            }
            s.append((char) data[i]);
            i++;
        }
        tag.setActionPool(actionPool);
        return i;
    }
}