#! /usr/bin/env python3
#
# Copyright © 2024 Collabora Ltd. and Red Hat Inc.
# SPDX-License-Identifier: MIT

# This script takes a list of Rust files, each of the form nvh_path_to_mod.rs
# and constructs a lib.rs which puts each of them in ::path::to::mod.

import argparse
import csv
import os
import sys

from mako import template

TEMPLATE_RS = template.Template(text="""\
// Copyright 2024 Red Hat Inc.
// SPDX-License-Identifier: MIT

// This file is generated by lat_rs_gen.py. DO NOT EDIT!
#![allow(unused_variables)]

const fn pred(has_pred: bool, a: u32, b: u32) -> u32 {
    if has_pred {
        a + b
    } else {
        b
    }
}

% for reg_file, cats in file_cats.items():
<% enum_name = to_camel(reg_file) + 'Latency' + sm.upper() %>
#[derive(PartialEq)]
pub enum ${enum_name} {
% for category in cats[0].header.cats:
   ${to_camel(category)},
% endfor
}

impl ${to_camel(reg_file)}Latency${sm.upper()} {
% for bigcat in cats:
    pub fn ${bigcat.header.latcat}(
        ${bigcat.header.cat0}: ${enum_name},
        ${bigcat.header.cat1}: ${enum_name},
        has_pred: bool
    ) -> u32 {
        use ${enum_name}::*;
        match ${bigcat.header.cat1} {
% for cat in bigcat.header.cats:
            ${to_camel(cat)} => match ${bigcat.header.cat0} {
            <% has_non = False %>
% for cat2 in bigcat.header.cats:
            % if bigcat.fields[loop.parent.index].flds[cat2].pred == True:
                ${to_camel(cat2)} => pred(
                    has_pred,
                    ${bigcat.fields[loop.parent.index].flds[cat2].value},
                    ${bigcat.fields[loop.parent.index].flds[cat2].pred_val}
                ),
            % elif bigcat.fields[loop.parent.index].flds[cat2].value != "none":
                ${to_camel(cat2)} => ${bigcat.fields[loop.parent.index].flds[cat2].value},
            % else:
                <% has_none = True %>
            % endif
% endfor
            % if has_none:
              _ => panic!("Illegal ${bigcat.header.cat0} value in ${bigcat.header.latcat} for ${to_camel(cat)}"),
            % endif
         }
% endfor
      }
   }
% endfor
}

% endfor
""")

## A mere convenience to convert snake_case to CamelCase. Numbers are prefixed
## with "_".
def to_camel(snake_str):
    result = ''.join(word.title() for word in snake_str.split('_'))
    return result if not result[0].isdigit() else '_' + result

def reader(csvfile):
    """Wrapper around csv.reader that skips comments and blanks."""
    # csv.reader actually reads the file one line at a time (it was designed to
    # open excel generated sheets), so hold the file until all of the lines are
    # read.
    with open(csvfile, 'r') as f:
        for line in csv.reader(f):
            if line and not line[0].startswith('#'):
                yield line

class Fld(object):
    def __init__(self, line):
        if "none" in line:
            self.valid = False
        else:
            self.valid = True
        self.pred = False
        if "+" in line:
            self.pred = True
            part = line.split("+")
            self.value = part[0]
            self.pred_val = part[1]
        elif " & sb" in line:
            self.scoreboard = True
            self.value = line.removesuffix(" & sb");
        else:
            self.scoreboard = False
            self.value = line.strip()

class Header(object):
    def __init__(self, line):
        self.latcat = line[0].strip()
        self.cats = line[1:]

        if self.latcat == "raw":
            self.cat0 = "writer"
            self.cat1 = "reader"
        elif self.latcat == "war":
            self.cat0 = "reader"
            self.cat1 = "writer"
        elif self.latcat == "waw":
            self.cat0 = "writer1"
            self.cat1 = "writer2"


class Fields(object):
    def __init__(self, header, line):
        self.fldcat = line[0].strip()
        self.flds = {}
        for index, cat in enumerate(header.cats):
            self.flds[cat] = Fld(line[index + 1])

class Category(object):
    def __init__(self, header, fields):
        self.header = header
        self.fields = fields

lattypes = ["reg", "ureg", "pred", "upred"]

def emit_cats(dirname, f, sm, lat):
    cats = []
    for index, cat in enumerate(["raw", "war", "waw"]):
        first_line = False
        fields = []
        for l in reader(dirname + "sm" + sm + "/" + lat + "_" + cat + ".csv"):
            if first_line == False:
                header = Header(l)
                first_line = True
            else:
                fields.append(Fields(header, l))
        cats.append(Category(header, fields))

    try:
            f.write(TEMPLATE_RS.render(lat=lat, sm=sm, cats=cats))
    except Exception:
        # In the event there's an error, this imports some helpers from mako
        # to print a useful stack trace and prints it, then exits with
        # status 1, if python is run with debug; otherwise it just raises
        # the exception
        import sys
        from mako import exceptions
        print(exceptions.text_error_template().render(), file=sys.stderr)
        sys.exit(1)

def emit_sm(dirname, f, sm):
    for lat in lattypes:
        emit_cats(dirname, f, sm, lat)

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('-p', '--import-path', required=True)
    parser.add_argument('--out-rs', required=True, help='Output Rust file.')
    parser.add_argument('--sm', help='SM', required=True)
    parser.add_argument('csv_files', metavar='FILE', nargs='*',
                        action='append',
                        help='Input CSV filename')

    args = parser.parse_args()

    sys.path.insert(0, args.import_path)
    import util

    file_cats = {}
    for csv_file in args.csv_files[0]:
        split = os.path.basename(csv_file).removesuffix('.csv').split('_')
        assert len(split) == 2
        reg_file = split[0]
        latcat = split[1]

        r = reader(csv_file)
        header = Header(next(r))
        fields = [Fields(header, l) for l in r]

        cat = Category(header, fields)
        file_cats.setdefault(reg_file, []).append(cat)

    environment = dict(
        sm=args.sm,
        file_cats=file_cats,
        to_camel=to_camel,
    )
    util.write_template_rs(args.out_rs, TEMPLATE_RS, environment)

if __name__ == '__main__':
    main()