from prompt_env9 import *
from LLM import *
from sre_constants import error
import random
import os
import json
import re
import copy
import numpy as np
import shutil
import time

import pdb
import sys
import re
import wandb
import warnings
import yaml
import time
from dotenv import load_dotenv
from tasks import load_task
from llm import load_llm

def path_constructor(loader, node):
    path_matcher = re.compile(r'\$\{([^}^{]+)\}')
    ''' Extract the matched value, expand env variable, and replace the match '''
    value = node.value
    match = path_matcher.match(value)
    env_var = match.group()[2:-1]
    return os.environ.get(env_var) + value[match.end():]


def load_config(cfg_path, args):
  path_matcher = re.compile(r'\$\{([^}^{]+)\}')
  yaml.add_implicit_resolver('!path', path_matcher)
  yaml.add_constructor('!path', path_constructor)
  with open(cfg_path, "r") as f:
    config = yaml.load(f, Loader=yaml.FullLoader)
  llm_config = config["llm"]
  agent_config = config["agent"]
  env_config = config["env"]
  run_config = config["run"]

  if args.log_path != '':
    run_config["log_path"] = args.log_path
  if args.project_name != '':
    run_config["project_name"] = args.project_name
  if args.baseline_dir != '':
    run_config["baseline_dir"] = args.baseline_dir

  run_config["wandb"] = args.wandb
  run_config["max_num_steps"] = args.max_num_steps

  return llm_config, agent_config, env_config, run_config


def check_log_paths_are_ready(log_dir, baseline_dir):
  if not os.path.exists(log_dir):
    os.makedirs(log_dir)

  if not os.path.exists(os.path.join(log_dir, "logs")):
    os.makedirs(os.path.join(log_dir, "logs"))

  if not os.path.exists(baseline_dir):
    os.makedirs(baseline_dir)

  if not os.path.exists(os.path.join(log_dir, 'all_results.txt')):
    with open(os.path.join(log_dir, 'all_results.txt'), "w") as f:
      f.write("")
      f.close()

  return True