import fxconv
import re
import os.path

from PIL import Image, ImageChops

def convert(input, output, params, target):
    recognized = True

    if params["custom-type"] == "level_map":
        o = convert_level_map(input, params)
    elif params["custom-type"] == "animation":
        o = convert_animation(input, params)
    elif params["custom-type"] == "aseprite-anim":
        o = convert_aseprite_anim(input, output, params)
    else:
        recognized = False

    if recognized:
        fxconv.elf(o, output, "_" + params["name"], **target)
        return 0

    return 1

def convert_level_map(input, params):
    RE_CATEGORY = re.compile(r'^\[([^\]]+)\]$', re.MULTILINE)

    with open(input, "r") as fp:
        ct = [x.strip() for x in RE_CATEGORY.split(fp.read())]
        ct = [x for x in ct if x != ""]
        assert(len(ct) % 2 == 0)

    ct = { ct[i]: ct[i+1] for i in range(0, len(ct), 2) }
    assert(ct.keys() == { "level", "base", "decor" })

    # Read tileset
    def tileset(desc):
        desc = [line.strip().split() for line in desc.split("\n")]
        assert(len(line) <= 16 for line in desc)

        tiles = dict()
        for y, line in enumerate(desc):
            for x, chara in enumerate(line):
                tiles[chara] = 16 * y + x

        return tiles

    base = tileset(ct["base"])
    decor = tileset(ct["decor"])

    level = [line.strip().split() for line in ct["level"].split("\n")]
    height = len(level)
    width = max(len(line) for line in level)

    tiles = bytearray(2 * width * height)
    for y, line in enumerate(level):
        for x, tile in enumerate(line):
            index = 2 * (width * y + x)
            tiles[index]   = base[tile[0]]
            tiles[index+1] = decor[tile[1]]

    o = bytes()
    o += fxconv.u16(width) + fxconv.u16(height)
    o += bytes(tiles)
    return o

def convert_animation(input, params):
    img = Image.open(input).convert("RGBA")

    fsize = [0,0]
    center = (0,0)
    durations = None
    frame_count = None
    next_anim = None

    # Read parameters

    if "frame_duration" in params:
        durations = [int(s) for s in params["frame_duration"].split(",")]
    else:
        durations = []

    if "frame_size" in params:
        fsize = list(map(int, params["frame_size"].split("x")))

    if "frame_height" in params:
        fsize[1] = int(params["frame_height"])
    elif not fsize[1]:
        fsize[1] = img.height

    if "frame_width" in params:
        fsize[0] = int(params["frame_width"])
    elif not fsize[0] and durations:
        fsize[0] = img.width // len(durations)
    else:
        raise fxconv.FxconvError("frame_width, frame_size, and frame_duration"+
            " are all unspecified - tell me more!")

    if "center" in params:
        center = tuple(map(int, params["center"].split(",")))
    else:
        center = (fsize[0] // 2, fsize[1] // 2)

    next_anim = params.get("next")

    name = params["name"]

    # Read and minimize frames

    grid = fxconv.Grid({ "width": fsize[0], "height": fsize[1] })
    frame_count = grid.size(img)

    bg = Image.new(img.mode, fsize, (0, 0, 0, 0))

    frames = []
    total_width = 0
    total_height = 0

    i = 0
    for rect in grid.iter(img):
        frame = img.crop(rect)
        nonzero = ImageChops.difference(frame, bg)
        bbox = nonzero.getbbox()
        if bbox:
            ox, oy = bbox[0], bbox[1]
            frame = frame.crop(bbox)
        else:
            ox, oy = 0, 0

        total_width += frame.width
        total_height = max(total_height, frame.height)

        frames.append((frame, ox, oy))
        i += 1

    # Create compact spritesheet and data arrays

    if total_width > 256:
        raise fxconv.FxconvError("uh total_width>256, call Lephe' x_x")
    sheet = Image.new("RGBA", (total_width, total_height))

    o = fxconv.ObjectData()
    sizeof_frame = 16
    x = 0

    for i, (frame, ox, oy) in enumerate(frames):
        o += fxconv.ref(f"{name}_sheet")
        o += fxconv.u8(x) + fxconv.u8(0)
        o += fxconv.u8(frame.width) + fxconv.u8(frame.height)
        o += fxconv.u8(center[0] - ox)
        o += fxconv.u8(center[1] - oy)
        o += fxconv.u16(durations[i])

        if i < len(frames) - 1:
            o += fxconv.ref(name, (i+1) * sizeof_frame)
        elif next_anim:
            o += fxconv.ref(next_anim)
        else:
            o += fxconv.u32(0)

        sheet.paste(frame, (x, 0))
        x += frame.width

    o += fxconv.sym(f"{name}_sheet")
    o += fxconv.convert_bopti_cg(sheet, {
        "profile": params.get("profile", "p8"),
    })

    return o

def convert_aseprite_anim(input, output, params):
    from aseprite import AsepriteFile

    #---
    # Perform checks for support on the Aseprite file
    #---

    with open(input, "rb") as fp:
        ase = AsepriteFile(fp.read())

    if ase.header.color_depth != 32:
        raise fxconv.FxconvError("Only RGBA supported yet, sorry x_x")

    # Find tags that give names to animations
    frame_tags = None
    for c in ase.frames[0].chunks:
        if c.chunk_type == 0x2018: # Tags
            frame_tags = c

    if frame_tags is None:
        raise fxconv.FxconvError("Found no frame tags")

    if len(ase.layers) != 1:
        raise fxconv.FxconvError("Only one layer supported yet, sorry x_x")

    # Find cel chunks of suitable types in each frame
    cel_chunks = []
    for i, f in enumerate(ase.frames):
        for c in f.chunks:
            if c.chunk_type == 0x2005: # Cel
                cel_chunks.append(c)
                break
        else:
            raise fxconv.FxconvError(f"Found no cels for frame #{i+1}")

        if cel_chunks[-1].cel_type not in [0,2]:
            raise fxconv.FxconvError(f"Cel chunk for frame #{i+1} is linked " +
                "or tilemap")

    #---
    # Print summary of animations
    #---

    print(f"{os.path.basename(input)} ({ase.header.width}x{ase.header.height})",
        f"has {len(frame_tags.tags)} animations:")
    for t in frame_tags.tags:
        name, from_, to = t['name'], t['from'], t['to']
        print(f"  '{name}': Frames {from_} to {to}", end="")
        durations = [ase.frames[i].frame_duration for i in range(from_, to+1)]
        print(" (" + ", ".join(f"{d} ms" for d in durations) + ")")

    #---
    # Generate PIL images for each frame
    #---

    pil_frames = []

    for i, f in enumerate(ase.frames):
        img = Image.new("RGBA", (ase.header.width, ase.header.height))
        c = cel_chunks[i]
        pixels = []
        offset = 0

        for y in range(c.data['height']):
            for x in range(c.data['width']):
                r, g, b, a = c.data['data'][offset:offset+4]
                pixels.append((r, g, b, a))
                offset += 4

        img.putdata(pixels)
        pil_frames.append(img)

    #---
    # Parse parameters
    #---

    if "center" in params:
        center = tuple(map(int, params["center"].split(",")))
    else:
        center = (ase.header.width // 2, ase.header.height // 2)

    if "next" in params:
        next_anim = [s.strip() for s in params["next"].split(",")]
        next_anim = dict([s.split("=", 1) for s in next_anim])
    else:
        next_anim = dict()

    #---
    # Generate compact sheets and object data for each animation
    #---

    bg = Image.new("RGBA", (ase.header.width, ase.header.height), (0, 0, 0, 0))

    # This will old all animations in sequence
    o = fxconv.ObjectData()
    sizeof_frame = 16

    for t in frame_tags.tags:
        name, from_, to = t['name'], t['from'], t['to']
        total_width = 0
        total_height = 0

        for i in range(from_, to+1):
            bbox = ImageChops.difference(pil_frames[i], bg).getbbox()
            if bbox:
                pil_frames[i] = (bbox[0], bbox[1], pil_frames[i].crop(bbox))
            else:
                pil_frames[i] = (0, 0, pil_frames[i])

            total_width += pil_frames[i][2].width
            total_height = max(total_height, pil_frames[i][2].height)


        # Define a new symbol for each animation
        symname = params["name"] + "_" + name
        s = fxconv.Structure()

        sheet = Image.new("RGBA", (total_width, total_height), (0, 0, 0, 0))
        x = 0

        for i in range(from_, to+1):
            s += fxconv.ref(f"{symname}_sheet")
            s += fxconv.u8(x) + fxconv.u8(0)
            s += fxconv.u8(pil_frames[i][2].width)
            s += fxconv.u8(pil_frames[i][2].height)
            s += fxconv.u8(center[0] - pil_frames[i][0])
            s += fxconv.u8(center[1] - pil_frames[i][1])
            s += fxconv.u16(ase.frames[i].frame_duration)

            if i < to:
                s += fxconv.ref(symname, (i-from_+1) * sizeof_frame)
            elif name in next_anim:
                s += fxconv.ref(params["name"] + "_" + next_anim[name])
            else:
                s += fxconv.u32(0)

            sheet.paste(pil_frames[i][2], (x, 0))
            x += pil_frames[i][2].width

        o += fxconv.sym(symname)
        o += s
        o += fxconv.sym(f"{symname}_sheet")
        o += fxconv.convert_bopti_cg(sheet, {
            "profile": params.get("profile", "p8"),
        })

    return o