#!/usr/bin/python3
#
# Script to maintain the project.json metadata file

import argparse
import json
import os
import re
import shutil
import sys

from dataclasses import dataclass, asdict

@dataclass
class ProjectMetadata:
    AssemblyName: str = ''
    AssemblyNumber: str = ''
    AssemblyTitle: str = ''
    Company: str = 'Asymworks, LLC'
    Designer: str = 'JPK'
    FabNumber: str = ''
    FabTitle: str = ''
    ProjectNumber: str = ''
    SchematicNumber: str = ''
    SchematicTitle: str = ''


#: Mapping from `ProjectMetadata` key to KiCad Text Variable name
TEXT_VARS = {
    'AssemblyName':     'ASSEMBLY_NAME',
    'AssemblyNumber':   'ASSEMBLY_NUMBER',
    'AssemblyTitle':    'DWG_TITLE_ASSY',
    'Company':          'COMPANY',
    'Designer':         'DESIGNER',
    'FabNumber':        'DWG_NUMBER_PCB',
    'FabTitle':         'DWG_TITLE_PCB',
    'ProjectNumber':    'PROJECT_CODE',
    'SchematicNumber':  'DWG_NUMBER_SCH',
    'SchematicTitle':   'DWG_TITLE_SCH',
}


def read_input(
    prompt: str,
    default: str | None = None,
    *,
    format_help: str | None = None,
    optional: bool = False,
    regex: str | re.Pattern | None = None,
    use_color: bool = True
) -> str | None:
    """Read an input string from the console, with defaults and retry."""
    prompt_str = f'{prompt} [{default or None}]'
    if use_color:
        prompt_str = f'{prompt} [\033[1;34m{default or None}\033[0m]'

    check_re = None
    if isinstance(regex, re.Pattern):
        check_re = regex
    elif isinstance(regex, str):
        check_re = re.compile(regex, re.I)

    error = None
    while True:
        error_str = f'\033[31m({error})\033[0m ' if error is not None else ''
        value = input(f'{error_str}{prompt_str}: ')
        if not value and not default and not optional:
            error = 'input required'
            continue

        if not value:
            value = default

        if check_re and not check_re.match(value):
            error = 'invalid format' if not format_help else format_help
            continue

        break

    if not value and optional:
        return None

    return value


def load_file(filename: str) -> ProjectMetadata:
    """Load the Project Metadata from a JSON File."""
    with open(filename, 'r') as f:
        data = json.load(f)

    return ProjectMetadata(**data)


def load_interactive(defaults: ProjectMetadata | None = None) -> ProjectMetadata:
    """Load the Project Metadata from an interactive console questionnaire."""
    md = defaults or ProjectMetadata()

    # Get the Company and Designer information first.
    md.Company = read_input('Company', md.Company)
    md.Designer = read_input('Designer', md.Designer)

    # Next get the Assembly Number and Name; this will allow us to
    # generate sensible defaults for the rest.
    md.AssemblyNumber = read_input('Assembly Number', md.AssemblyNumber)
    md.AssemblyName = read_input('Assembly Name', md.AssemblyName)

    if not md.AssemblyTitle:
        md.AssemblyTitle = f'Assembly, {md.AssemblyName}'

    if not md.FabNumber:
        md.FabNumber = f'P{md.AssemblyNumber[1:]}'
    if not md.FabTitle:
        md.FabTitle = f'PCB Fabrication, {md.AssemblyName}'

    if not md.SchematicNumber:
        md.SchematicNumber = f'S{md.AssemblyNumber[1:]}'
    if not md.SchematicTitle:
        md.SchematicTitle = f'Schematic, {md.AssemblyName}'

    md.AssemblyTitle = read_input('Assembly Drawing Title', md.AssemblyTitle)
    md.FabNumber = read_input('Fabrication Drawing Number', md.FabNumber)
    md.FabTitle = read_input('Fabrication Drawing Title', md.FabTitle)
    md.SchematicNumber = read_input('Schematic Drawing Number', md.SchematicNumber)
    md.SchematicTitle = read_input('Schematic Drawing Title', md.SchematicTitle)

    # Finally load the Project Number
    if not md.ProjectNumber:
        md.ProjectNumber = f'P{md.AssemblyNumber[1:3]}'

    md.ProjectNumber = read_input(
        'Project Number',
        md.ProjectNumber,
        format_help='Format as P##',
        regex=r'^P[0-9]{2}$',
    )

    return md


def run_init(filename: str) -> int:
    """Run the `init` subcommand to initialize the Metadata."""
    md = None
    if os.path.isfile(filename):
        md = load_file(filename)

    md = load_interactive(md)
    with open(filename, 'w') as f:
        json.dump(asdict(md), f, indent=2)

    return 0


def run_print(filename: str, variable: str) -> int:
    """Run the `print` subcommand to echo Metadata to `stdout`."""
    try:
        md = load_file(filename)
    except FileNotFoundError:
        sys.stderr.write(f'{filename} not found\n')
        return 1
    except TypeError as e:
        sys.stderr.write(f'{filename} has unexpected data\n  {str(e)}\n')
        return 1

    var_lcase = variable.lower()
    for md_key, txt_var in TEXT_VARS.items():
        if var_lcase != md_key.lower() and var_lcase != txt_var.lower():
            continue

        var_found = True
        if not hasattr(md, md_key):
            sys.stderr.write(f'Internal error: {md_key} not defined in metadata')
            return 2

        sys.stdout.write(getattr(md, md_key))
        return 0

    sys.stderr.write(f'{variable} is not a valid Metadata variable\n')
    return 1

def run_update(filename: str, kicad_project: str) -> int:
    """Run the `update` subcommand to update the KiCad Project."""
    try:
        md = load_file(filename)
    except FileNotFoundError:
        sys.stderr.write(f'{filename} not found\n')
        return 1
    except TypeError as e:
        sys.stderr.write(f'{filename} has unexpected data\n  {str(e)}\n')
        return 1

    kicad_filepath = os.path.abspath(kicad_project)
    if not os.path.isfile(kicad_filepath):
        sys.stderr.write(f'{kicad_filepath} does not exist\n')
        return 1

    if os.path.splitext(kicad_filepath)[1] != '.kicad_pro':
        sys.stderr.write(f'{kicad_filepath} is not a KiCad Project file\n')
        return 1

    kicad_data = None
    try:
        with open(kicad_filepath, 'r') as f:
            kicad_data = json.load(f)

        if 'text_variables' not in kicad_data:
            sys.stderr.write(f'Warning: "text-variables" is not set in KiCad Project')

        for md_key, txt_var in TEXT_VARS.items():
            if not hasattr(md, md_key):
                sys.stderr.write(f'Internal error: {md_key} not defined in metadata')
                return 2

            if txt_var not in kicad_data['text_variables']:
                sys.stderr.write(f'Warning: creating KiCad text variable {txt_var}')

            kicad_data['text_variables'][txt_var] = getattr(md, md_key)

        shutil.copy(kicad_filepath, f'{kicad_filepath}-bak')
        with open(kicad_filepath, 'w') as f:
            json.dump(kicad_data, f, indent=2)

    except FileNotFoundError:
        sys.stderr.write(f'{kicad_filepath} does not exist\n')
        return 1

    return 0


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script to maintain the project.json metadata file')
    parser.add_argument('-f', '--filename', type=str, default='project.json', help='Project Metadata filename (default is "project.json")')
    
    subparsers = parser.add_subparsers(dest='command', metavar='command', help='subcommands')

    init_parser = subparsers.add_parser('init', help='Initialize the Project Metadata')
    print_parser = subparsers.add_parser('print', help='Print Metadata values')
    print_parser.add_argument('variable', help='Metadata variable name to print')
    update_parser = subparsers.add_parser('update', help='Update the KiCad Project with the Metadata')
    update_parser.add_argument('kicad_project', help='KiCad project to update with the Metadata')

    # parser.add_argument('-i', '--interactive', action='store_true', help='Run in interactive mode to prompt for each data field.')
    # parser.add_argument('-u', '--update', nargs=1, metavar='KICAD-PROJECT', help='Write the text variables in the KiCad Project file')

    args = parser.parse_args()
    filepath = os.path.abspath(args.filename)

    if args.command == 'init':
        sys.exit(run_init(filepath))

    elif args.command == 'print':
        sys.exit(run_print(filepath, args.variable))

    elif args.command == 'update':
        sys.exit(run_update(filepath, args.kicad_project))

    sys.stderr.write(f'Unknown command: {args.command}\n')
    sys.exit(1)
