main.py

import os, sys
import json, jsonschema
from pathlib import Path
import subprocess
import argparse

from jslt_parser.jslt_parser import parser
from jslt_parser.pvs_walker import PVSConverter

def log_message(message: str, indent_level: int = 0):
    """Utility function to log messages with consistent formatting and script name.
    Args:
        message: The message to log
        indent_level: Number of spaces to indent the message (default: 0)
    """
    indent = " " * indent_level
    print(f"[{sys.argv[0]}] {indent}{message}")

class PathManager:
    def __init__(self, pvs_location: str, pvs_output: str, instance_name: str):
        self.pvs_location = Path(pvs_location).resolve()
        self.pvs_output = Path(pvs_output).resolve()
        self.instance_name = instance_name
        
        self.dockerized = os.getenv("FHIRFLY_DOCKER", None) == "true"
        # Validate PVS location
        if not self.pvs_location.exists():
            raise ValueError(f"PVS location does not exist: {self.pvs_location}")
            
        # Create output directory if it doesn't exist
        self.pvs_output.mkdir(parents=True, exist_ok=True)
        
        # Get the project root directory (where main.py is located)
        self.project_root = Path(__file__).parent.resolve()
        
        # Calculate the actual jsltlib directory path
        self.jsltlib_dir = (self.project_root / "jslt_pvs/jsltlib").resolve()
        
        # Validate jsltlib exists
        if not self.jsltlib_dir.exists():
            raise ValueError(f"jsltlib not found at {self.jsltlib_dir}")
            
        # Calculate relative path from pvs_output to jsltlib
        self.jsltlib_path = os.path.relpath(self.jsltlib_dir, self.pvs_output)

        # Get the local PVS environment setup script
        if not self.dockerized:
            self.pvs_env_script = (self.pvs_location / "setup-env.sh").resolve()
            if not self.pvs_env_script.exists():
                raise ValueError(f"PVS environment setup script not found at {self.pvs_env_script}")

    @property
    def pvs_executable(self) -> Path:
        
        if self.dockerized:
            return self.pvs_location / "pvs"
        else:
            return self.pvs_location / "pvs-local"
    @property
    def pvs_theory_file(self) -> Path:
        return self.pvs_output / f"{self.instance_name}.pvs"
        
    @property
    def pvs_strategies_file(self) -> Path:
        return self.pvs_output / "pvs-strategies"
        
    @property
    def pvs_batch_file(self) -> Path:
        return self.pvs_output / f"{self.instance_name}.el"
        
    @property
    def pvs_log_file(self) -> Path:
        return self.pvs_output / f"{self.instance_name}.log"
        
    @property
    def jsltlib_relative_path(self) -> str:
        """Returns the static jsltlib path"""
        return self.jsltlib_path
        
    def get_pvs_command(self) -> str:
        """Returns the command to run PVS with the local environment"""
        if self.dockerized:
            return f"cd {self.pvs_output} && {self.pvs_executable} -batch -v 1 -q -l {self.instance_name}.el"
        else:
            # Source the environment setup script and then run PVS
            return f"source {self.pvs_env_script} && cd {self.pvs_output} && {self.pvs_executable} -batch -v 1 -q -l {self.instance_name}.el"

    def setup_pvs_environment(self):
        """Sets up the PVS environment variables for subprocess calls"""
        env = os.environ.copy()
        
        if self.dockerized:
            return env
        # Read the setup-env.sh script and extract environment variables
        try:
            result = subprocess.run(
                f"source {self.pvs_env_script} && env",
                shell=True,
                capture_output=True,
                text=True,
                executable="/bin/bash"
            )
            
            for line in result.stdout.splitlines():
                if "=" in line:
                    key, value = line.split("=", 1)
                    env[key] = value
                    
        except subprocess.CalledProcessError as e:
            log_message(f"Failed to source PVS environment: {e}")
            raise
            
        return env

class PVSManager:
    def __init__(self, name):
        self.parser = parser
        self.walker = PVSConverter()
        
        self.name = name
        self.theories = ["jslt"]
        self.requires = None
        self.ensures = None
        self.program = None

        self.registered = False
        self.ran_pvs = False

        self.path_manager = None

    def base_to_jsondata(self, base_data):
        if isinstance(base_data, int):
            return f"jnumeral(juint64({base_data}))"
        elif isinstance(base_data, str):
            return f'jstr("{base_data}")'
        elif isinstance(base_data, list):
            raise Exception("Unimplemented - array types!")
            return f"jarray({base_data})"
        else: # object case
            return base_data
        
    def json_to_jsondata(self, obj):
        s = ",".join([f'jpair("{k}", {self.base_to_jsondata(v)})' for (k, v) in obj.items()])
        return "jdict(" + s + ")"

        
    def jslt_to_pvs(self, program):
        return self.walker.walk(node=program)

    def requires_to_pvs(self):
        if not self.registered:
            raise Exception("Must call register(input_schema, transform, output_schema) before outputting `requires`!")
        
        if self.requires.get_node_type_str() != "null":
            log_message("Found `requires` declaration in JSLT, translating constraint to PVS")
            return self.jslt_to_pvs(self.requires)

        log_message("Did not find `requires` keyword in JSLT, using dsl_true")
        return "dsl_true"

    def ensures_to_pvs(self):
        if not self.registered:
            raise Exception("Must call register(input_schema, transform, output_schema) before outputting `ensures`!")
        if self.ensures.get_node_type_str() != "null":
            log_message("Found `ensures` declaration in JSLT, translating constraint to PVS")
            return self.jslt_to_pvs(self.ensures)

        log_message("Did not find `ensures` keyword in JSLT, using dsl_true")
        return "dsl_true"


    def __typecheck(self, schema, key):
        type = schema['properties'][key]['type']
        if type == "string":
            return ("jstr?", "string")
        elif type == "integer":
            return ("jnumeral?", "number")
        else:
            raise ValueError(f"JSON base types must only be string/integer, not {type}")
    
    def schema_to_pvs(self, schema, metavariable="obj"):
        if not self.registered:
            raise Exception("Must call register(input_schema, transform, output_schema) before outputting schema constraints!")
        s = json.loads(schema)
        log_message(f"Translating schema {s['title']}")

        l = len(s['required'])

        constraint = f""" LET j = contents({metavariable}) IN j`length = {l} AND % object must have (at least) {l} elements """
        for (i, key) in enumerate(s['required']):
            (pvs_type, pp_type) = self.__typecheck(s, key)
            constraint += f"""
  % {{ "{key}": [{pp_type} type] }}
  (LET pair = j`seq({i}) IN key(pair) = "{key}" AND {pvs_type}(value(pair))) { "AND" if i != len(s['required'])-1 else "" } """
            
        return constraint
                  
    def register(self, input_schema, transform, output_schema, instance_name, pvs_location, pvs_output, constraints):
        self.registered = True
        self.input_schema = self.schema_to_pvs(input_schema)
        self.output_schema = self.schema_to_pvs(output_schema)

        p = parser.parse(transform)

        c_p = parser.parse(constraints)

        self.requires = c_p.requires.expr
        self.ensures = c_p.ensures.expr
        
        self.transform = self.jslt_to_pvs(p.expr)
        
        # Initialize path manager
        self.path_manager = PathManager(pvs_location, pvs_output, instance_name)
        
        log_message(f"PVS output path = {self.path_manager.pvs_output}")
        
        return

    def emit_pvs_verif(self):
        req = self.requires_to_pvs()
        ens = self.ensures_to_pvs()
        
        return f'''
{self.name}_verif: THEORY
BEGIN

  jsltlib: LIBRARY = "{self.path_manager.jsltlib_relative_path}/"

  IMPORTING {",".join(['jsltlib@' + theory for theory in self.theories])}

  input_compliant?(obj: jsondata): bool =
  {self.input_schema}

  % predicate subtyping - the type of JSON objects that are input compliant
  % mathematically (and in PVS), {{ j: jsondata | input_compliant?(j) }}
  input_compliant: TYPE = (input_compliant?)

  output_compliant?(obj: jsondata): bool =
  {self.output_schema}

  % predicate subtyping - the type of JSON objects that are output compliant
  % mathematically (and in PVS), {{ j: jsondata | output_compliant?(j) }}
  output_compliant: TYPE = (output_compliant?)

  % standard library/helper utilities
  dsl_true: MACRO dsl_expr = Apply(LT, (# length := 2, seq := [: Constant(jnumeral(juint64(0))), Constant(jnumeral(juint64(1))) :] #))
  dsl_false: MACRO dsl_expr = Apply(GT, (# length := 2, seq := [: Constant(jnumeral(juint64(0))), Constant(jnumeral(juint64(1))) :] #))

  requires(pvs_input_obj: jsondata): dsl_expr =
  { "dsl_true" if self.requires == None else req }

  ensures(pvs_input_obj, pvs_output_obj: jsondata): dsl_expr =
  { "dsl_true" if self.ensures == None else ens }

  transform(pvs_input_obj: (jdict?)): dsl_expr = {self.transform}
        
  correctness: THEOREM
    FORALL (obj: input_compliant):
      { "TRUE" if self.requires == None else "OUT_2(interpret(requires(obj))(obj))" } IMPLIES
      LET new_obj = OUT_1(interpret(transform(obj))(obj)) IN
      output_compliant?(new_obj) AND
      { "TRUE" if self.ensures == None else "OUT_2(interpret(ensures(obj, new_obj))(obj))" }

END {self.name}_verif'''
        
    def emit_pvs_exec(self):
        return f'''
{self.name}_exec: THEORY
BEGIN
  IMPORTING {self.name}_verif
        
  jsltlib: LIBRARY = "{self.path_manager.jsltlib_relative_path}/"
  IMPORTING {",".join(['jsltlib@' + theory for theory in self.theories])}

  transform(pvs_input_obj: (jdict?)): dsl_expr = { self.transform }

  transform_{self.name}: [(jdict?) -> jsondata] =
    LAMBDA (obj: (jdict?)): get_obj(raw_eval(transform(obj))(obj))

  requires_{self.name}: [(jdict?) -> boolean] =
    LAMBDA (obj: (jdict?)): b(raw_eval(requires(obj))(obj))

  ensures_{self.name}: [[(jdict?), (jdict?)] -> boolean] =
     LAMBDA (obj: [(jdict?), (jdict?)]): b(raw_eval(ensures(proj_1(obj), proj_2(obj)))(proj_2(obj)))

END {self.name}_exec

        '''

    def emit_pvs(self):
        if not self.registered:
            raise Exception("Must call register(input_schema, transform, output_schema) before outputting PVS!")
        
        verif = self.emit_pvs_verif()
        exec = self.emit_pvs_exec()
        return verif + "\n\n\n\n" + exec

    def emit_pvs_to_file(self, filename):
        if not self.registered:
            raise Exception("Must call register(input_schema, transform, output_schema) before outputting PVS file!")
        str = self.emit_pvs()
        with open(filename, "w+") as f:
            f.write(str)

    def emit_pvs_strategies_file(self):
        return """
(defstep jslt-verify-interpret (&optional timeout)
  (apply
   (then@
    (rewrite-msg-off)
    (SKEEP)
    (TYPEPRED "obj")
    (EXPAND "input_compliant?")
    (GROUND)
    (EXPAND "interpret")
    (EXPAND "transform")
    (ASSERT)
    (GRIND))
   :time? t
   :timeout timeout)
   "Documentation: proof strategy for JSLT verification
    &optional timeout - positive integer denoting number of seconds before proof skipped"
   "")
        """

    def emit_pvs_batch_file(self, instance_name):
        return f"""
(pvs-validate
 "{instance_name}.log"
 "./"
 (pvs-message "Running proofs for {instance_name}")
 (let ((current-prefix-arg t))
   (prove-formulas-theory "{instance_name}_verif" "(jslt-verify-interpret)")))
        """

    def run_pvs(self, verify=False):
        if self.ran_pvs:
            raise Exception("Running PVS again?")
        self.ran_pvs = True

        log_message("%%%%%%%%%%%%%%%% SUMMARY")
        log_message("Runtime functions generated:")
        log_message(f"         transform_{self.name} : [jsondata -> jsondata]", 10)
        
        # Only show requires/ensures if they exist and are not None
        if self.requires is not None and self.requires.get_node_type_str() != "null":
            log_message(f"          requires_{self.name} : [jsondata -> boolean]", 10)
        if self.ensures is not None and self.ensures.get_node_type_str() != "null":
            log_message(f"          ensures_{self.name}  : [[jsondata, jsondata] -> boolean]", 10)

        command = self.path_manager.get_pvs_command()

        if not verify:
            log_message("Run the following command for verification:")
            log_message(f"                  {command}")

    def validate_pvs_run(self, instance_name):
        if not self.ran_pvs:
            raise Exception("Must run PVS before validating PVS run!")
            
        log_file = self.path_manager.pvs_log_file
        
        if not log_file.exists():
            raise Exception("Can't find {log_file}!")

        flag = False
        with open(log_file, 'r') as f:
            for line in f:
                if f"correctness proved" in line:
                    log_message("Proved!")
                    flag = True
                elif f"correctness unproved" in line:
                    break

        if not flag:
            log_message("Verification failed!")
        else:
            log_message("Verification succeeded!")

        


def run(input_schema,
        transform,
        output_schema,
        instance_name,
        pvs_location,
        pvs_output,
        constraints_file,
        verify=False):
    
    # Read all input files
    input_files = {
        'input_schema': Path(input_schema),
        'output_schema': Path(output_schema),
        'transform': Path(transform),
        'constraints': Path(constraints_file)
    }
    
    file_contents = {}
    for name, path in input_files.items():
        if not path.exists():
            raise FileNotFoundError(f"Could not find {name} file at {path}")
        log_message(f"Reading {path} for {name}")
        file_contents[name] = path.read_text()

    log_message("Initializing PVS manager")
    pvs = PVSManager(instance_name)
    log_message("Done initializing PVS manager")

    log_message("Registering inputs with PVS manager")
    pvs.register(input_schema=file_contents['input_schema'],
                 transform=file_contents['transform'],
                 output_schema=file_contents['output_schema'],
                 instance_name=instance_name,
                 pvs_location=pvs_location,
                 pvs_output=pvs_output,
                 constraints=file_contents['constraints'])
    log_message("Done registering inputs with PVS manager")

    # Write PVS theory file
    theory_file = pvs.path_manager.pvs_theory_file
    log_message(f"Writing to {theory_file}")
    theory_file.write_text(pvs.emit_pvs())
    log_message(f"Done writing PVS theories to {theory_file}")

    # Write PVS strategies file
    strategies_file = pvs.path_manager.pvs_strategies_file
    log_message(f"Populating PVS strategy file: {strategies_file}")
    strategies_file.write_text(pvs.emit_pvs_strategies_file())
    log_message(f"Done writing PVS strategy file @ {strategies_file}")

    # Write PVS batch file
    batch_file = pvs.path_manager.pvs_batch_file
    log_message(f"Writing PVS batch script command: {batch_file}")
    batch_file.write_text(pvs.emit_pvs_batch_file(instance_name))
    log_message(f"Done writing PVS batch script @ {batch_file}")

    pvs.run_pvs(verify)

    # If verify flag is set, run the PVS command
    if verify:
        log_message("Running PVS verification...")
        command = pvs.path_manager.get_pvs_command()
        
        # Run the command with the proper PVS environment
        try:
            env = pvs.path_manager.setup_pvs_environment()
            result = subprocess.run(
                command,
                shell=True,
                check=True,
                capture_output=True,
                text=True,
                env=env,
                executable="/bin/bash"
            )
            log_message("PVS verification completed")
            
            # Validate the run
            pvs.validate_pvs_run(instance_name)
            
        except subprocess.CalledProcessError as e:
            log_message("PVS verification failed with error:")
            log_message(e.stderr)
            sys.exit(1)

def main():
    parser = argparse.ArgumentParser(
        prog="fhirfly",
        description="FHIR-Fly: PVS/Yices-backed platform for generating and verifying JSONSchema-to-JSONSchema transformations"
    )
    parser.add_argument("--input_schema",
                        help="Input JSONSchema describing form and fields of the input data format (.json)",
                        required=True)
    parser.add_argument("--transform",
                        help="Input JSLT transformation converting instances of `input_schema` to instances of `output_schema` (.jslt)",
                        required=True)
    parser.add_argument("--output_schema",
                        help="Input JSONSchema describing form and fields of the output data format (.json)",
                        required=True)
    parser.add_argument("--instance_name",
                        default="demo",
                        help="Optional instance name (default=`demo`)")
    parser.add_argument("--pvs_location",
                        help="Location of PVS installation (folder containing setup-env.sh)",
                        required=True)
    parser.add_argument("--pvs_output",
                        default="results",
                        help='Optional location of .pvs file (+ auxiliary files - strategy efile, batch mode script) to be generated (default="results")')
    parser.add_argument("--constraint_file",
                        help='Optional location of "requires"/"ensures" constraints')
    parser.add_argument("--verify",
                        help="Automatically run PVS verification after generating files",
                        action="store_true")

    args = parser.parse_args()

    constraints_file = args.constraint_file if args.constraint_file else args.transform

    run(input_schema=args.input_schema,
        transform=args.transform,
        output_schema=args.output_schema,
        instance_name=args.instance_name,
        pvs_location=args.pvs_location,
        pvs_output=args.pvs_output,
        constraints_file=constraints_file,
        verify=args.verify)
    
if __name__ == '__main__':
    global script_name
    script_name = sys.argv[0]
    main()