# Copyright (c) 2024-present, Royal Bank of Canada.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
#

import os

from utils import wrap_code
import pdb
import base64
import re

def encode_image(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')
          

ONE_SHOT_INIT_PROMPT_TEMPLATE = """You are given a natural language description of a planning problem in the domain {target_domain_name} along with one problem instance in PDDL format and an image observation of the initial setup. Your task is to generate a PDDL domain for the target domain {target_domain_name} that is equivalent to its natural language description and is compatible with the provided problem instance. 

Starting from a PDDL domain template, you are allowed to modify the template using the following two python function interfaces:

```python
add_or_update_predicates(predicates: List[str])
modify_action(action_name: str, new_preconditions: List[str], new_effects: List[str])
```

An example of above functions applied to an example PDDL domain template is as follows:

{context_shot_example}

Target Domain Description:
{target_domain_nl}

Target Problem PDDL:
{target_problem_pddl}

Now, your task is to complete the following PDDL template by generating necessary predicates and action preconditions and effects:

Target PDDL Template:
{target_domain_template_pddl}

You must never modify action parameters, and you are only allowed to use the following two function interfaces to modify the template.

"""

BLOCKS_WORLD_EXAMPLE = """Example Domain Description:
```markdown
The robot has four actions: pickup, putdown, stack, and unstack. The domain assumes a world where there are a set of blocks that can be stacked on top of each other, an arm that can hold one block at a time, and a table where blocks can be placed.
The actions defined in this domain include:
pickup: allows the arm to pick up a block from the table if it is clear and the arm is empty. After the pickup action, the arm will be holding the block, and the block will no longer be on the table or clear.
putdown: allows the arm to put down a block on the table if it is holding a block. After the putdown action, the arm will be empty, and the block will be on the table and clear.
stack: allows the arm to stack a block on top of another block if the arm is holding the top block and the bottom block is clear. After the stack action, the arm will be empty, the top block will be on top of the bottom block, and the bottom block will no longer be clear.
unstack: allows the arm to unstack a block from on top of another block if the arm is empty and the top block is clear. After the unstack action, the arm will be holding the top block, the top block will no longer be on top of the bottom block, and the bottom block will be clear.
```

Example Problem PDDL:
```pddl
(define (problem BW-rand-5)
  (:domain blocksworld-4ops)
  (:objects b1 b2 b3 b4 b5 )
  (:init
    (arm-empty)
    (on b1 b4)
    (on b2 b5)
    (on b3 b2)
    (on-table b4)
    (on b5 b1)
    (clear b3)
  )
  (:goal
    (and
      (on b4 b3)
    )
  )
)
```

Example PDDL Template:
```pddl
(define (domain blocksworld-4ops)
  (:requirements :strips)
  (:predicates)

  (:action pickup
    :parameters (?ob)
    :precondition ()
    :effect ()
  )

  (:action putdown
    :parameters (?ob)
    :precondition ()
    :effect ()
  )

  (:action stack
    :parameters (?ob ?underob)
    :precondition ()
    :effect ()
  )

  (:action unstack
    :parameters (?ob ?underob)
    :precondition ()
    :effect ()
  )
)
```

Example Completion:
```python
add_or_update_predicates([
    '(clear ?x)',
    '(on-table ?x)',
    '(arm-empty)',
    '(holding ?x)',
    '(on ?x ?y)'
])
modify_action('pickup', [
    "(clear ?ob)",
    "(on-table ?ob)",
    "(arm-empty)"
], [
    "(holding ?ob)",
    "(not (clear ?ob))",
    "(not (on-table ?ob))",
    "(not (arm-empty))"
])

modify_action('putdown', [
    "(holding ?ob)"
], [
    "(clear ?ob)",
    "(arm-empty)",
    "(on-table ?ob)",
    "(not (holding ?ob))"
])

modify_action('stack', [
    "(clear ?underob)",
    "(holding ?ob)"
], [
    "(arm-empty)",
    "(clear ?ob)",
    "(on ?ob ?underob)",
    "(not (clear ?underob))",
    "(not (holding ?ob))"
])

modify_action('unstack', [
    "(on ?ob ?underob)",
    "(clear ?ob)",
    "(arm-empty)"
], [
    "(holding ?ob)",
    "(clear ?underob)",
    "(not (on ?ob ?underob))",
    "(not (clear ?ob))",
    "(not (arm-empty))"
])
```
"""


MAZE_EXAMPLE = """Example Domain Description:
```markdown
The robot has four actions: move-up, move-down, move-left, and move-right. This domain models an agent navigating in a grid world. The goal is to reach a target location. The world consists of a set of discrete positions (e.g., grid cells). The position representation is (row, column) representation. For example, (pos-4-3) represents the position in fourth row and third column, and (pos-4-4) is at right of (pos-4-3). Between each two positions, the directional links between them design how they can move between each other. For example, the agent can move-left from pos-3-2 to pos-3-1 because pos-3-1 is at left to pos-3-2, but cannot move-left from pos-3-2 to pos-3-3. 
```

Example Problem PDDL:
```pddl

(define (problem maze) (:domain maze)
  (:objects
        pos-1-1 - position
        pos-1-2 - position
        pos-1-3 - position
        pos-2-1 - position
        pos-2-2 - position
        pos-2-3 - position
        pos-3-1 - position
        pos-3-2 - position
        pos-3-3 - position
  )
  (:init 
    (at pos-2-1)
    (is-goal pos-3-3)
    (move-dir-left pos-1-2 pos-1-1)
    (move-dir-left pos-1-3 pos-1-2)
    (move-dir-left pos-2-2 pos-2-1)
    (move-dir-left pos-2-3 pos-2-2)
    (move-dir-left pos-3-2 pos-3-1)
    (move-dir-left pos-3-3 pos-3-2)
    (move-dir-right pos-1-1 pos-1-2)
    (move-dir-right pos-1-2 pos-1-3)
    (move-dir-right pos-2-1 pos-2-2)
    (move-dir-right pos-2-2 pos-2-3)
    (move-dir-right pos-3-1 pos-3-2)
    (move-dir-right pos-3-2 pos-3-3)
    (move-dir-up pos-2-1 pos-1-1)
    (move-dir-up pos-2-2 pos-1-2)
    (move-dir-up pos-2-3 pos-1-3)
    (move-dir-up pos-3-1 pos-2-1)
    (move-dir-up pos-3-2 pos-2-2)
    (move-dir-up pos-3-3 pos-2-3)
    (move-dir-down pos-1-1 pos-2-1)
    (move-dir-down pos-1-2 pos-2-2)
    (move-dir-down pos-1-3 pos-2-3)
    (move-dir-down pos-2-1 pos-3-1)
    (move-dir-down pos-2-2 pos-3-2)
    (move-dir-down pos-2-3 pos-3-3)
  )
  (:goal (at pos-3-2))
)
```

Example PDDL Template:
```pddl
(define (problem maze)
    (:domain maze)
    (:objects )
    (:init )
    (:goal (and ))
)
```

Example Completion:
```python
add_or_update_predicates([
    '(move-dir-up ?x ?y)',
    '(move-dir-down ?x ?y)',
    '(move-dir-left ?x ?y)',
    '(move-dir-right ?x ?y)',
    '(at ?x)',
    '(is-goal ?x)'
])
modify_action('move-up', [
    "(at ?from)",
    "(move-dir-up ?from ?to)"
], [
    "(not (at ?from))",
    "(not (clear ?to))",
    "(at ?to)",
])

modify_action('move-down', [
    "(at ?from)",
    "(move-dir-down ?from ?to)"
], [
    "(not (at ?from))",
    "(at ?to)",
])

modify_action('move-left', [
    "(at ?from)",
    "(move-dir-left ?from ?to)"
], [
    "(not (at ?from))",
    "(at ?to)",
])

modify_action('move-right', [
    "(at ?from)",
    "(move-dir-right ?from ?to)"
], [
    "(not (at ?from))",
    "(at ?to)",
])
```
"""
ACTION_LEVEL_INIT_PROMPT_TEMPLATE = """You are given a natural language description of a planning problem in the domain {target_domain_name} along with one problem instance in PDDL format and an image observation of the initial setup. Your task is to generate a PDDL domain for the target domain {target_domain_name} that is equivalent to its natural language description and is compatible with the provided problem instance. The following is an example of a valid PDDL domain:

{context_domain_pddl}

Natural language description of the domain:

{target_domain_nl}

PDDL description of the problem:

{target_problem_pddl}

Now, your task is to complete the following PDDL template by generating necessary predicates and action preconditions and effects:

{target_domain_template_pddl}

You must never modify action parameters, and you are only allowed to use the following two function interfaces to modify the template.

```python
add_or_update_predicates(predicates: List[str])
modify_action(action_name: str, new_preconditions: List[str], new_effects: List[str])
```

An example of above functions is as follows:

```python
add_or_update_predicates(['(is-robot ?x)'])
modify_action('move', ['(is-robot ?x)'], ['(is-robot ?y)'])
```
"""

# Problem Translation

PROBLEM_TRANSLATION_SYSTEM_MESSAGE = """You are a helpful assistant, skilled in producing Planning Domain Definition Language (PDDL) code of problems given a natural language description of a planning problem."""
with open(os.path.join(os.path.dirname(__file__), os.pardir, 'data_my/domains/blocksworld/domain.pddl'), 'r') as f:
    BLOCKS_WORLD_DOMAIN_PDDL_WRAPPED = wrap_code(f.read(), lang='pddl')
with open(os.path.join(os.path.dirname(__file__), os.pardir, 'data_my/domains/blocksworld/domain.nl'), 'r') as f:
    BLOCKS_WORLD_DOMAIN_NL_WRAPPED = wrap_code(f.read(), lang='markdown')
    
PROBLEM_TRANSLATION_BLOCKS_WORLD_INPUT = """Domain Description:
{domain_nl_wrapped}
{maybe_context_domain_pddl}
Domain PDDL Template:
```pddl
(define (domain blocksworld)
  (:requirements :strips)
  (:predicates)

  (:action pickup
    :parameters (?ob)
    :precondition ()
    :effect ()
  )

  (:action putdown
    :parameters (?ob)
    :precondition ()
    :effect ()
  )

  (:action stack
    :parameters (?ob ?underob)
    :precondition ()
    :effect ()
  )

  (:action unstack
    :parameters (?ob ?underob)
    :precondition ()
    :effect ()
  )
)
Problem Description:
```markdown
You are tasked with manipulating a set of blocks using a robotic arm that can perform four actions: pickup, putdown, stack, and unstack.

Initially:
- The robotic arm is empty.
- Block b_green is on block b_blue.
- Block b_orange is on block b_red.
- Block b_yellow is on block b_orange, and it is clear.
- Block b_blue is on the table.
- Block b_red is on block b_green.

Your goal is to achieve the following configuration:
- Block b_blue must be stacked on top of block b_yellow.
```

Problem PDDL Template:
```pddl
(define (problem BW)
    (:domain blocksworld)
    (:objects )
    (:init )
    (:goal (and ))
)
```
"""

PROBLEM_TRANSLATION_BLOCKS_WORLD_OUTPUT = """Problem PDDL:
```markdown
(define (problem BW)
  (:domain blocksworld)
  (:objects b_green b_orange b_yellow b_blue b_red )
  (:init
    ; The robotic arm is empty.
    (arm-empty)
    ; Block b_green is on block b_blue.
    (on b_green b_blue)
    ; Block b_orange is on block b_red.
    (on b_orange b_red)
    ; Block b_yellow is on block b_orange, and it is clear.
    (on b_yellow b_orange)
    (clear b_yellow)
    ; Block b_blue is on the table.
    (on-table b_blue)
    ; Block b_red is on block b_green.
    (on b_red b_green)
  )
  (:goal
    (and
      ; Block b_blue must be stacked on top of block b_yellow.
      (on b_blue b_yellow)
    )
  )
)
```
"""

with open(os.path.join(os.path.dirname(__file__), os.pardir, 'data_my/domains/maze/domain.pddl'), 'r') as f:
    MAZE_DOMAIN_PDDL_WRAPPED = wrap_code(f.read(), lang='pddl')
with open(os.path.join(os.path.dirname(__file__), os.pardir, 'data_my/domains/maze/domain.nl'), 'r') as f:
    MAZE_DOMAIN_NL_WRAPPED = wrap_code(f.read(), lang='markdown')
PROBLEM_TRANSLATION_MAZE_INPUT = """Domain Description:
{domain_nl_wrapped}
{maybe_context_domain_pddl}
Domain PDDL Template:
```pddl
(define (domain maze)
  (:requirements :strips)
  (:predicates)

  (:action move-up
    :parameters (?from ?to)
    :precondition ()
    :effect ()
  )

  (:action move-down
    :parameters (?from ?to)
    :precondition ()
    :effect ()
  )

  (:action move-left
    :parameters (?from ?to)
    :precondition ()
    :effect ()
  )

  (:action move-right
    :parameters (?ob ?underob)
    :precondition ()
    :effect ()
  )
)
```
Problem Description:
```markdown
You are tasked with manipulating an agent to reach the goal without colliding into the wall. The position representation is (row, column) representation. For example, (pos-4-3) represents the position in fourth row and third column. The left upper corner is (pos-1-1). You can perform four actions: move-up, move-down, move-left, and move-right.

Initially:
- The agent is at (pos-2-1).
- The goal is at (pos-3-3).

Your goal is to achieve the following configuration:
- The agent is at (pos-3-3)
```

Problem PDDL Template:
```pddl
(define (problem maze)
    (:domain maze)
    (:objects )
    (:init )
    (:goal (and ))
)
```
"""
PROBLEM_TRANSLATION_MAZE_OUTPUT = """Problem PDDL:
```markdown

(define (problem maze) (:domain maze)
  (:objects
        pos-1-1 - position
        pos-1-2 - position
        pos-1-3 - position
        pos-2-1 - position
        pos-2-2 - position
        pos-2-3 - position
        pos-3-1 - position
        pos-3-2 - position
        pos-3-3 - position
  )
  (:init 
    (at pos-2-1)
    (is-goal pos-3-3)
    (move-dir-left pos-1-2 pos-1-1)
    (move-dir-left pos-1-3 pos-1-2)
    (move-dir-left pos-2-2 pos-2-1)
    (move-dir-left pos-2-3 pos-2-2)
    (move-dir-left pos-3-2 pos-3-1)
    (move-dir-left pos-3-3 pos-3-2)
    (move-dir-right pos-1-1 pos-1-2)
    (move-dir-right pos-1-2 pos-1-3)
    (move-dir-right pos-2-1 pos-2-2)
    (move-dir-right pos-2-2 pos-2-3)
    (move-dir-right pos-3-1 pos-3-2)
    (move-dir-right pos-3-2 pos-3-3)
    (move-dir-up pos-2-1 pos-1-1)
    (move-dir-up pos-2-2 pos-1-2)
    (move-dir-up pos-2-3 pos-1-3)
    (move-dir-up pos-3-1 pos-2-1)
    (move-dir-up pos-3-2 pos-2-2)
    (move-dir-up pos-3-3 pos-2-3)
    (move-dir-down pos-1-1 pos-2-1)
    (move-dir-down pos-1-2 pos-2-2)
    (move-dir-down pos-1-3 pos-2-3)
    (move-dir-down pos-2-1 pos-3-1)
    (move-dir-down pos-2-2 pos-3-2)
    (move-dir-down pos-2-3 pos-3-3)
  )
  (:goal (at pos-3-2))
)

```
"""
PROBLEM_TRANSLATION_USER_INPUT_TEMPLATE = """Domain Description:
{target_domain_nl}
{maybe_target_domain_pddl}
Domain PDDL Template:
{target_domain_template_pddl}

Problem Description:
{target_problem_nl}

Problem PDDL Template:
{target_problem_template_pddl}
"""


def get_problem_translation_messages(
        target_domain_nl, target_problem_nl, target_problem_template_pddl, target_domain_template_pddl, target_domain_pddl=None, context_img_path=None, problem_img_path=None, multimodal=False
):
    if target_domain_pddl is not None:
      if 'maze' not in context_img_path:
        maybe_context_domain_pddl = f"\nDomain PDDL:\n{BLOCKS_WORLD_DOMAIN_PDDL_WRAPPED}"
      else:
        maybe_context_domain_pddl = f"\nDomain PDDL:\n{MAZE_DOMAIN_PDDL_WRAPPED}"
      maybe_target_domain_pddl = f"\nDomain PDDL:\n{target_domain_pddl}"
    else:
        maybe_context_domain_pddl = ""
        maybe_target_domain_pddl = ""
      
    if 'maze' not in context_img_path:
      input_prompt = PROBLEM_TRANSLATION_BLOCKS_WORLD_INPUT
      output_prompt = PROBLEM_TRANSLATION_BLOCKS_WORLD_OUTPUT
      domain_nl_wrapped=BLOCKS_WORLD_DOMAIN_NL_WRAPPED
    else:
      input_prompt = PROBLEM_TRANSLATION_MAZE_INPUT
      output_prompt = PROBLEM_TRANSLATION_MAZE_OUTPUT
      domain_nl_wrapped=MAZE_DOMAIN_NL_WRAPPED
      # input_prompt = (re.sub(r'pos-(\d+)-(\d+)', r'pos-\2-\1', PROBLEM_TRANSLATION_MAZE_INPUT)).replace('column', 'column_no').replace('row', 'column').replace('column_no', 'row').replace('third row and fourth column', 'forth row and third column')
      
    # pdb.set_trace()
    context_initial_image = encode_image(context_img_path)
    target_initial_image = encode_image(problem_img_path)
    if multimodal:
      # pdb.set_trace()
      return PROBLEM_TRANSLATION_SYSTEM_MESSAGE, [
          {'role': 'user',
          'content': 'Given a natural language description of a planning problem and an image observation of the initial scenario, your task is to generate a complete PDDL problem instance that is equivalent to its natural language description and image observation, contain all necessary predicates as described in the description, and is thorough and complete. Consider all predicates that are relevant and neccesary. Do not omit anything or leave for me to fill in.'},
          {'role': 'assistant',
          'content': 'Sure, please provide the natural language description and image observation of the planning problem, and I will help you generate the corresponding PDDL problem instance.'},
          {'role': 'user', 'content': [{"type": "text", "text": input_prompt.format(
              domain_nl_wrapped=domain_nl_wrapped,
              maybe_context_domain_pddl=maybe_context_domain_pddl
          )}, {"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{context_initial_image}"}}]},
          {'role': 'assistant', 'content': output_prompt},
      ], (PROBLEM_TRANSLATION_USER_INPUT_TEMPLATE.format(
          target_domain_nl=target_domain_nl,
          maybe_target_domain_pddl=maybe_target_domain_pddl,
          target_domain_template_pddl=target_domain_template_pddl,
          target_problem_nl=target_problem_nl,
          target_problem_template_pddl=target_problem_template_pddl,
      ), target_initial_image)
    else:
      return PROBLEM_TRANSLATION_SYSTEM_MESSAGE, [
          {'role': 'user',
          'content': 'Given a natural language description of a planning problem of the initial scenario, your task is to generate a complete PDDL problem instance that is equivalent to its natural language description and is thorough and complete. Do not omit anything or leave for me to fill in.'},
          {'role': 'assistant',
          'content': 'Sure, please provide the natural language description of the planning problem, and I will help you generate the corresponding complete and full PDDL problem instance.'},
          {'role': 'user', 'content': input_prompt.format(
              domain_nl_wrapped=domain_nl_wrapped,
              maybe_context_domain_pddl=maybe_context_domain_pddl
          )},
          {'role': 'assistant', 'content': output_prompt},
      ], (PROBLEM_TRANSLATION_USER_INPUT_TEMPLATE.format(
          target_domain_nl=target_domain_nl,
          maybe_target_domain_pddl=maybe_target_domain_pddl,
          target_problem_nl=target_problem_nl,
          target_problem_template_pddl=target_problem_template_pddl,
      ), target_initial_image)


# Domain Translation

DOMAIN_TRANSLATION_SYSTEM_MESSAGE = """You are a helpful assistant, skilled in producing Planning Domain Definition Language (PDDL) code of domains given a natural language description of a planning domain and an image observation of the initial scenario. You always wrap your code in the appropriate markdown or PDDL syntax."""
DOMAIN_TRANSLATION_CONTEXT_DOMAIN_INPUT = """Domain Description:
{context_domain_nl}

Problem Description:
{context_problem_nl}

Domain PDDL Template:
{context_domain_template_pddl}
"""
DOMAIN_TRANSLATION_CONTEXT_DOMAIN_OUTPUT = """Domain PDDL:
{context_domain_pddl}
"""
DOMAIN_TRANSLATION_USER_INPUT_TEMPLATE = """Domain Description:
{target_domain_nl}

Problem Description:
{target_problem_nl}

Domain PDDL Template:
{target_domain_template_pddl}
"""


def get_domain_translation_messages(
        context_domain_nl, context_problem_nl, context_domain_template_pddl, context_domain_pddl,
        target_domain_nl, target_problem_nl, target_domain_template_pddl
):
    return DOMAIN_TRANSLATION_SYSTEM_MESSAGE, [
        {'role': 'user',
         'content': 'Given a natural language description of a planning domain in PDDL format as well as domain template and an example problem description, your task is to generate a complete PDDL domain that is equivalent to its natural language description and is thorough and complete.'},
        {'role': 'assistant',
         'content': 'Sure, I would be happy to help you generate a PDDL domain based on the natural language description and domain template you provide. Please share the natural language description of the planning domain, the domain template, and the example problem description, and I will assist you in creating the corresponding PDDL domain.'},
        {'role': 'user', 'content': DOMAIN_TRANSLATION_CONTEXT_DOMAIN_INPUT.format(
            context_domain_nl=context_domain_nl,
            context_problem_nl=context_problem_nl,
            context_domain_template_pddl=context_domain_template_pddl
        )},
        {'role': 'assistant', 'content': DOMAIN_TRANSLATION_CONTEXT_DOMAIN_OUTPUT.format(
            context_domain_pddl=context_domain_pddl
        )},
    ], DOMAIN_TRANSLATION_USER_INPUT_TEMPLATE.format(
        target_domain_nl=target_domain_nl,
        target_problem_nl=target_problem_nl,
        target_domain_template_pddl=target_domain_template_pddl
    )


ONE_SHOT_PROBLEM_TRANSLATION_PROMPT = """Your task is to generate PDDL code for problems given a natural language description of a planning problem, an image observation of the initial problem setup, and a PDDL template for the problem instance.
To achieve this task, you are also given a natural language description of the domain, the domain PDDL, and a PDDL problem instance.
Always wrap your code in the appropriate PDDL syntax.

Domain Description:
{domain_nl}

Domain PDDL:
{domain_pddl}

Q:
Problem Description:
{context_problem_nl}
Problem PDDL Template:
{context_problem_template_pddl}

A:
{context_problem_pddl}

Now, please provide a PDDL code for the following problem:

Q:
Problem Description:
{target_problem_nl}
Problem PDDL Template:
{target_problem_template_pddl}

A:
"""


FILES_UPDATE_SYSTEM_MESSAGE = """You are a helpful assistant, skilled in producing Planning Domain Definition Language (PDDL) code of problems and domains given a natural language description of a planning domain and an image observation of the initial scenario. You always wrap your code in the appropriate markdown or PDDL syntax."""
TWO_FILES_UPDATE_PROMPT = """Here are your generated PDDL problem and domain. Please reason about the issue with your generated problem file or generated code for domain generation.
Problem Description:
{target_problem_nl}

Domain Description:
{target_domain_nl}

Domain PDDL Template:
{target_domain_template}

Generated Problem PDDL:
{problem_pddl}

Generated Domain PDDL:
{domain_pddl}

Execution Feedback:
{execution_feedback}

In your response, please 
1. reason about if you want to update your problem file and domain file to fix the issue by answering
Q1: List all predicates. Describe how they are related to objects and actions and how they should be updated. 
Q2: List all actions in domain. Describe, explain, and compare action's definition in given natural language description with its inputs, preconditions, and effects in domain file one by one. No two actions should have same preconditions and effects. If not, please correct in domain file.
Q3: List all objects in problem. 
Q4: Check the :objects section. If fixed 'typing' exists, list the type of each object and if it is EXPLICITLY written as format [object_name - type] in problem file. Explicit typing means the exact string - <type> must appear after each object name. Do not assume implicit typing. If any object is listed without a type, then the problem file is incorrect and must be modified.
Q5: List all init states in problem. Predicates used in domain should be initialized in problem file. And they should describe every aspect of the problem, and be enough to execute actions. When defining directional links of a grid, the init states should exhaustively enumerate all adjacent positions for every applicable cell in this grid, from first row and column to last row and column, no matter it is valid movement or contains obstacles or not. If not, please correct in problem file.
Q6: List all goals in problem. 
Q7: Based on previous answers, summarize the errors one by one, and explain which file(s) should be updated and which parts of them should be updated in detail.

2. update the problem file to be correct and complete according to your analysis, if needed (Your modified PDDL will be directly used, so return only complete and valid PDDL: no comments, TODOs, placeholders, or omissions); 

3. generate the new code for domain generation according to your analysis, if needed. 

If you want to modify the new domain file, you are only allowed to use the following two function interfaces to modify the template.

```python
add_or_update_predicates(predicates: List[str])
modify_action(action_name: str, new_preconditions: List[str], new_effects: List[str])
```

An example of above functions is as follows:

```python
add_or_update_predicates(['(is-robot ?x)'])
modify_action('move', ['(is-robot ?x)'], ['(is-robot ?y)'])
```

Please give your answer strictly in the following format:

Problem file and Domain file update reasoning: []
New problem file: [fill in N/A if not needed]
New domain file: [fill in N/A if not needed]"
"""


def get_two_files_modification_messages(
        target_domain_nl, target_problem_nl, target_domain_template, problem_pddl, domain_pddl, execution_feedback, problem_img, multimodal
):
    if multimodal:
      return FILES_UPDATE_SYSTEM_MESSAGE, [
          {'role': 'user',
          'content': 'Given a natural language description of a planning problem and an image observation of the initial scenario, you have generated a PDDL problem file and a PDDL domain file correspondingly. Now these two files are not completely correct in that it cannot execute some actions that should be executable. Your task is to, based on some feedback from executing the action, reason about what is problematic within your generated files, and fix the incorrect file or files.'},
          {'role': 'assistant',
          'content': 'Sure, I would be happy to check my generated files and fix the incorrect file(s) based on execution feedback. Please share the natural language description, PDDL problem and PDDL domain I created and the execution feedback.'}
      ], (TWO_FILES_UPDATE_PROMPT.format(
              target_problem_nl=target_problem_nl,
              target_domain_nl=target_domain_nl,
              target_domain_template=target_domain_template,
              problem_pddl=problem_pddl,
              domain_pddl=domain_pddl,
              execution_feedback=execution_feedback
          ), problem_img)