#!/usr/bin/env python3
"""Prepare a local test atlas from an existing Java client JAR (no downloads).

Keeps original pixel colors; smaller sprites use integer nearest-neighbor
scaling. Animated sprites use their first declared frame. Face materials are
projected onto Shacraft's existing geometry, not a replacement model renderer.
Requires Pillow. Generated assets must not be redistributed without permission.
"""
import argparse
from functools import lru_cache
import hashlib
from io import BytesIO
import json
import math
from pathlib import Path
import shutil
import zipfile

from PIL import Image

ROOT = Path(__file__).resolve().parents[1]
PREFIX = 'assets/minecraft/'
DIRECTIONS = ['east', 'west', 'up', 'down', 'south', 'north']
NORMALS = [(1, 0, 0), (-1, 0, 0), (0, 1, 0), (0, -1, 0), (0, 0, 1), (0, 0, -1)]


def norm(name):
    return name.removeprefix('minecraft:')


def projection(direction, p):
    x, y, z = p
    return [(1-z, 1-y), (z, 1-y), (x, z), (x, 1-z), (x, 1-y), (1-x, 1-y)][direction]


def rotate(p, x, y):
    # Java blockstate rotations are clockwise around the block center.
    x, y = math.radians(-x), math.radians(-y)
    a, b, c = p
    b, c = b*math.cos(x)-c*math.sin(x), b*math.sin(x)+c*math.cos(x)
    return (a*math.cos(y)+c*math.sin(y), b, -a*math.sin(y)+c*math.cos(y))


def condition(when, props):
    if 'OR' in when:
        return any(condition(c, props) for c in when['OR'])
    if 'AND' in when:
        return all(condition(c, props) for c in when['AND'])
    return all(props.get(k) in str(v).split('|') for k, v in when.items())


class Compiler:
    def __init__(self, jar, images):
        self.jar, self.images = jar, images

    @lru_cache(None)
    def read(self, path):
        return json.loads(self.jar.read(PREFIX + path))

    @lru_cache(None)
    def model(self, name):
        data = self.read('models/' + norm(name) + '.json')
        parent = data.get('parent')
        base = self.model(parent) if parent and 'builtin/' not in parent else {}
        return {**base, **data, 'textures': {**base.get('textures', {}), **data.get('textures', {})}}

    def resolve(self, value, textures):
        seen = set()
        while isinstance(value, str) and value.startswith('#') and value not in seen:
            seen.add(value)
            value = textures.get(value[1:])
        if isinstance(value, dict):
            value = value.get('sprite')
        if isinstance(value, str) and norm(value).startswith('block/'):
            name = norm(value).removeprefix('block/')
            if name in self.images:
                return name
        return None

    @lru_cache(None)
    def model_faces(self, name, rx=0, ry=0, uvlock=False):
        model = self.model(name)
        textures = model.get('textures', {})
        result, areas = [None]*6, [-1]*6
        for element in model.get('elements', []):
            lo = [v/16 for v in element['from']]
            hi = [v/16 for v in element['to']]
            for direction, data in element.get('faces', {}).items():
                texture = self.resolve(data.get('texture'), textures)
                if not texture or texture == 'grass_block_side_overlay':
                    continue  # The opaque grass side already includes its edge.
                face = DIRECTIONS.index(direction)
                normal = tuple(round(v) for v in rotate(NORMALS[face], rx, ry))
                dest = NORMALS.index(normal)
                axes = [i for i, v in enumerate(NORMALS[face]) if not v]
                area = math.prod(abs(hi[i]-lo[i]) for i in axes)
                if area <= areas[dest]:
                    continue
                areas[dest] = area
                corners = [projection(face, p) for p in [lo, hi]]
                low = [min(c[i] for c in corners) for i in range(2)]
                high = [max(c[i] for c in corners) for i in range(2)]
                rect = [v/16 for v in data.get('uv', [low[0]*16, low[1]*16, high[0]*16, high[1]*16])]

                def uv_at(world):
                    # Inverse rotation is the transpose of the orthogonal basis.
                    offset = [v-.5 for v in world]
                    basis = [rotate(axis, rx, ry) for axis in [(1, 0, 0), (0, 1, 0), (0, 0, 1)]]
                    local = [.5 + sum(a*b for a, b in zip(offset, axis)) for axis in basis]
                    uv = projection(face, local)
                    if uvlock:
                        return projection(dest, world)
                    uv = [(uv[i]-low[i])/max(high[i]-low[i], 1e-6) for i in range(2)]
                    for _ in range(data.get('rotation', 0)//90):
                        uv = [uv[1], 1-uv[0]]
                    return [rect[i] + uv[i]*(rect[i+2]-rect[i]) for i in range(2)]

                zero = uv_at([0, 0, 0])
                unit = [uv_at(p) for p in [[1, 0, 0], [0, 1, 0], [0, 0, 1]]]
                transform = [round(v, 6) for i in range(2) for v in [*(p[i]-zero[i] for p in unit), zero[i]]]
                result[dest] = {'texture': texture, 'uv': transform, 'tinted': 'tintindex' in data}
        particle = self.resolve(textures.get('particle'), textures)
        fallback = next((v for v in result if v), None)
        for i, face in enumerate(result):
            if face is None:
                texture = particle or (fallback and fallback['texture'])
                result[i] = {'texture': texture, 'tinted': bool(fallback and fallback['tinted'])} if texture else None
        return result

    def faces(self, block, props):
        definition = self.read('blockstates/' + norm(block) + '.json')
        models = []
        for selector, model in definition.get('variants', {}).items():
            when = dict(part.split('=', 1) for part in selector.split(',') if part)
            if condition(when, props):
                models.append(model[0] if isinstance(model, list) else model)
                break
        for part in definition.get('multipart', []):
            if condition(part.get('when', {}), props):
                model = part['apply']
                models.append(model[0] if isinstance(model, list) else model)
        faces = [None]*6
        for model in models:
            incoming = self.model_faces(model['model'], model.get('x', 0), model.get('y', 0), model.get('uvlock', False))
            faces = [a or b for a, b in zip(faces, incoming)]
        result = []
        for face in faces:
            if not face:
                result.append(None)
                continue
            face = face.copy()
            tinted = face.pop('tinted', False)
            texture = face['texture']
            if tinted:
                tint = [0.58, 0.8, 0.34]
                if 'leaves' in block or 'vine' in block:
                    tint = [0.46, 0.7, 0.28]
                if 'spruce' in block:
                    tint = [0.38, 0.60, 0.38]
                if block == 'minecraft:redstone_wire':
                    power = int(props.get('power', '0'))/15
                    tint = [.3+.7*power, max(0, power*power*.7-.5), 0]
                face['tint'] = tint
            if block in ('minecraft:water', 'minecraft:bubble_column'):
                face['tint'] = [0.25, 0.46, 0.9]
            alpha = self.images[texture].getchannel('A').getextrema()
            face['cutout'] = alpha[0] == 0 and alpha[1] == 255
            result.append(face)
        return result


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument('--jar', type=Path, required=True)
    parser.add_argument('--states', type=Path, required=True, help='Matching generated/reports/blocks.json')
    parser.add_argument('--output', type=Path, required=True)
    args = parser.parse_args()
    output = args.output.resolve()
    if output.exists():
        parser.error('Output already exists; choose a new directory.')
    jar = zipfile.ZipFile(args.jar)
    images, animated, native_sizes = {}, [], {}
    for path in sorted(jar.namelist()):
        if not path.startswith(PREFIX + 'textures/block/') or not path.endswith('.png'):
            continue
        name = Path(path).stem
        image = Image.open(BytesIO(jar.read(path))).convert('RGBA')
        metadata = json.loads(jar.read(path + '.mcmeta')).get('animation', {}) if path + '.mcmeta' in jar.namelist() else {}
        width = metadata.get('width', min(image.size))
        height = metadata.get('height', width)
        frames = metadata.get('frames', [0])
        first = frames[0] if frames else 0
        index = first.get('index', 0) if isinstance(first, dict) else first
        x, y = index % (image.width//width)*width, index // (image.width//width)*height
        frame = image.crop((x, y, x+width, y+height))
        native_sizes[name] = [width, height]
        if metadata or image.size != frame.size:
            animated.append(name)
        images[name] = frame.resize((32, 32), Image.Resampling.NEAREST)
    compiler = Compiler(jar, images)
    report = json.loads(args.states.read_text())
    sets, lookup, states, defaults, uncovered = [], {}, {}, {}, []
    for block, definition in report.items():
        for state in definition['states']:
            props = state.get('properties', {})
            try:
                faces = compiler.faces(block, props)
            except KeyError:
                faces = [None]*6
            key = json.dumps(faces, separators=(',', ':'), sort_keys=True)
            if key not in lookup:
                lookup[key] = len(sets)
                sets.append(faces)
            canonical = block + ('['+','.join(f'{k}={v}' for k, v in sorted(props.items()))+']' if props else '')
            states[canonical] = lookup[key]
            if state.get('default'):
                defaults[block] = lookup[key]
                if not any(faces):
                    uncovered.append(block)
    output.mkdir(parents=True)
    for base in ('base', 'trampoline'):
        shutil.copytree(ROOT/'packages'/base, output/base)
    package = output/'vanilla'
    (package/'client').mkdir(parents=True)
    columns = 40
    rows = math.ceil(len(images)/columns)
    atlas = Image.new('RGBA', (columns*32, rows*32))
    for i, image in enumerate(images.values()):
        atlas.paste(image, (i % columns*32, i//columns*32))
    atlas.save(package/'client/atlas.png', optimize=True)
    descriptor = {'schema': 1, 'texture_pack': {
        'name': 'Minecraft 26.2 Original — Local Test', 'pixel_size': 32,
        'atlas': {'path': 'client/atlas.png', 'columns': columns, 'rows': rows},
        'textures': [{'name': name} for name in images],
        'grass_overlay': {'base': 'grass_block_side', 'overlay': 'grass_block_side_overlay'},
        'block_faces': {'sets': sets, 'states': states, 'defaults': defaults},
    }}
    (package/'client/style.json').write_text(json.dumps(descriptor, separators=(',', ':'))+'\n')
    resources = []
    for path, role in [('client/atlas.png', 'texture'), ('client/style.json', 'client-style')]:
        data = (package/path).read_bytes()
        assert len(data) < 16*1024*1024
        resources.append({'path': path, 'scope': 'client', 'role': role, 'size': len(data), 'sha256': hashlib.sha256(data).hexdigest()})
    base = json.loads((output/'base/manifest.json').read_text())
    manifest = {'schema': 1, 'id': 'shacraft.vanilla.local', 'version': '1.0.1',
                'license': 'Minecraft assets copyright Mojang/Microsoft. Local testing only; no redistribution rights granted.',
                'dependencies': [{'id': base['id'], 'version': base['version']}],
                'capabilities': ['client.texture', 'client.style'], 'resources': resources}
    (package/'manifest.json').write_text(json.dumps(manifest, indent=2)+'\n')
    receipt = {'textures': len(images), 'states': len(states), 'face_sets': len(sets),
               'source_sha256': hashlib.sha256(args.jar.read_bytes()).hexdigest(),
               'native_sizes': dict((str(size), list(native_sizes.values()).count(size)) for size in native_sizes.values()),
               'animated_first_frame': animated, 'untextured_default_blocks': uncovered,
               'atlas_bytes': (package/'client/atlas.png').stat().st_size,
               'mapping_bytes': (package/'client/style.json').stat().st_size,
               'limitations': ['Static first animation frame', 'Existing engine geometry; complex multipart and entity-rendered blocks are approximated', 'Fixed foliage tint rather than biome colors']}
    (output/'receipt.json').write_text(json.dumps(receipt, indent=2)+'\n')
    print(json.dumps(receipt, indent=2))


if __name__ == '__main__':
    main()