#!/usr/bin/env python3

import sys
import os
import os.path
import subprocess
import tempfile
import re
import itertools
import hashlib
import shutil
import argparse
import codecs
import json
import multiprocessing
import errno
from functools import partial

class Paths():
    def __init__(self, spirv_cross, glslang, spirv_as, spirv_val, spirv_opt):
        self.spirv_cross = spirv_cross
        self.glslang = glslang
        self.spirv_as = spirv_as
        self.spirv_val = spirv_val
        self.spirv_opt = spirv_opt

def remove_file(path):
    #print('Removing file:', path)
    os.remove(path)

def create_temporary(suff = ''):
    f, path = tempfile.mkstemp(suffix = suff)
    os.close(f)
    #print('Creating temporary:', path)
    return path

def parse_stats(stats):
    m = re.search('([0-9]+) work registers', stats)
    registers = int(m.group(1)) if m else 0

    m = re.search('([0-9]+) uniform registers', stats)
    uniform_regs = int(m.group(1)) if m else 0

    m_list = re.findall('(-?[0-9]+)\s+(-?[0-9]+)\s+(-?[0-9]+)', stats)
    alu_short = float(m_list[1][0]) if m_list else 0
    ls_short = float(m_list[1][1]) if m_list else 0
    tex_short = float(m_list[1][2]) if m_list else 0
    alu_long = float(m_list[2][0]) if m_list else 0
    ls_long = float(m_list[2][1]) if m_list else 0
    tex_long = float(m_list[2][2]) if m_list else 0

    return (registers, uniform_regs, alu_short, ls_short, tex_short, alu_long, ls_long, tex_long)

def get_shader_type(shader):
    _, ext = os.path.splitext(shader)
    if ext == '.vert':
        return '--vertex'
    elif ext == '.frag':
        return '--fragment'
    elif ext == '.comp':
        return '--compute'
    elif ext == '.tesc':
        return '--tessellation_control'
    elif ext == '.tese':
        return '--tessellation_evaluation'
    elif ext == '.geom':
        return '--geometry'
    else:
        return ''

def get_shader_stats(shader):
    path = create_temporary()

    p = subprocess.Popen(['malisc', get_shader_type(shader), '--core', 'Mali-T760', '-V', shader], stdout = subprocess.PIPE, stderr = subprocess.PIPE)
    stdout, stderr = p.communicate()
    remove_file(path)

    if p.returncode != 0:
        print(stderr.decode('utf-8'))
        raise OSError('malisc failed')
    p.wait()

    returned = stdout.decode('utf-8')
    return parse_stats(returned)

def print_msl_compiler_version():
    try:
        subprocess.check_call(['xcrun', '--sdk', 'iphoneos', 'metal', '--version'])
        print('...are the Metal compiler characteristics.\n')   # display after so xcrun FNF is silent
    except OSError as e:
        if (e.errno != errno.ENOENT):    # Ignore xcrun not found error
            raise

def path_to_msl_standard(shader):
    if '.ios.' in shader:
        if '.msl2.' in shader:
            return '-std=ios-metal2.0'
        elif '.msl21.' in shader:
            return '-std=ios-metal2.1'
        elif '.msl11.' in shader:
            return '-std=ios-metal1.1'
        elif '.msl10.' in shader:
            return '-std=ios-metal1.0'
        else:
            return '-std=ios-metal1.2'
    else:
        if '.msl2.' in shader:
            return '-std=macos-metal2.0'
        elif '.msl21.' in shader:
            return '-std=macos-metal2.1'
        elif '.msl11.' in shader:
            return '-std=macos-metal1.1'
        else:
            return '-std=macos-metal1.2'

def path_to_msl_standard_cli(shader):
    if '.msl2.' in shader:
        return '20000'
    elif '.msl21.' in shader:
        return '20100'
    elif '.msl11.' in shader:
        return '10100'
    else:
        return '10200'

def validate_shader_msl(shader, opt):
    msl_path = reference_path(shader[0], shader[1], opt)
    try:
        if '.ios.' in msl_path:
            msl_os = 'iphoneos'
        else:
            msl_os = 'macosx'
        subprocess.check_call(['xcrun', '--sdk', msl_os, 'metal', '-x', 'metal', path_to_msl_standard(msl_path), '-Werror', '-Wno-unused-variable', msl_path])
        print('Compiled Metal shader: ' + msl_path)   # display after so xcrun FNF is silent
    except OSError as oe:
        if (oe.errno != errno.ENOENT):   # Ignore xcrun not found error
            raise
    except subprocess.CalledProcessError:
        print('Error compiling Metal shader: ' + msl_path)
        raise RuntimeError('Failed to compile Metal shader')

def cross_compile_msl(shader, spirv, opt, iterations, paths):
    spirv_path = create_temporary()
    msl_path = create_temporary(os.path.basename(shader))

    spirv_cmd = [paths.spirv_as, '-o', spirv_path, shader]
    if '.preserve.' in shader:
        spirv_cmd.append('--preserve-numeric-ids')

    if spirv:
        subprocess.check_call(spirv_cmd)
    else:
        subprocess.check_call([paths.glslang, '--target-env', 'vulkan1.1', '-V', '-o', spirv_path, shader])

    if opt:
        subprocess.check_call([paths.spirv_opt, '--skip-validation', '-O', '-o', spirv_path, spirv_path])

    spirv_cross_path = paths.spirv_cross

    msl_args = [spirv_cross_path, '--entry', 'main', '--output', msl_path, spirv_path, '--msl', '--iterations', str(iterations)]
    msl_args.append('--msl-version')
    msl_args.append(path_to_msl_standard_cli(shader))
    if '.swizzle.' in shader:
        msl_args.append('--msl-swizzle-texture-samples')
    if '.ios.' in shader:
        msl_args.append('--msl-ios')
    if '.pad-fragment.' in shader:
        msl_args.append('--msl-pad-fragment-output')
    if '.capture.' in shader:
        msl_args.append('--msl-capture-output')
    if '.domain.' in shader:
        msl_args.append('--msl-domain-lower-left')
    if '.argument.' in shader:
        msl_args.append('--msl-argument-buffers')
    if '.discrete.' in shader:
        # Arbitrary for testing purposes.
        msl_args.append('--msl-discrete-descriptor-set')
        msl_args.append('2')
        msl_args.append('--msl-discrete-descriptor-set')
        msl_args.append('3')

    subprocess.check_call(msl_args)

    if not shader_is_invalid_spirv(msl_path):
        subprocess.check_call([paths.spirv_val, '--target-env', 'vulkan1.1', spirv_path])

    return (spirv_path, msl_path)

def shader_model_hlsl(shader):
    if '.vert' in shader:
        if '.sm30.' in shader:
            return '-Tvs_3_0'
        else:
            return '-Tvs_5_1'
    elif '.frag' in shader:
        if '.sm30.' in shader:
            return '-Tps_3_0'
        else:
            return '-Tps_5_1'
    elif '.comp' in shader:
        return '-Tcs_5_1'
    else:
        return None

def shader_to_win_path(shader):
    # It's (very) convenient to be able to run HLSL testing in wine on Unix-likes, so support that.
    try:
        with subprocess.Popen(['winepath', '-w', shader], stdout = subprocess.PIPE, stderr = subprocess.PIPE) as f:
            stdout_data, stderr_data = f.communicate()
            return stdout_data.decode('utf-8')
    except OSError as oe:
        if (oe.errno != errno.ENOENT): # Ignore not found errors
            return shader
    except subprocess.CalledProcessError:
        raise

    return shader

ignore_fxc = False
def validate_shader_hlsl(shader, force_no_external_validation, paths):
    subprocess.check_call([paths.glslang, '-e', 'main', '-D', '--target-env', 'vulkan1.1', '-V', shader])
    is_no_fxc = '.nofxc.' in shader