import os

from intercode.envs import (
  BashEnv, IntercodeEnv, AGENT_OBS, REWARD
)
from typing import Dict, Tuple

class AgentEnv(BashEnv):
    """Gym environmnet for Agent-bench"""
    name = "ic_agent"

    def __init__(self, image_name: str, **kwargs):
        IntercodeEnv.__init__(self, image_name, **kwargs)
        self.token = os.environ.get("GITHUB_TOKEN", "git")
    
    def reset_container(self) -> None:
        self.workdir = "/"
        folders = self.container.exec_run(self.clean_cmd('ls'), workdir="/").output.decode("utf-8")
        self.agent_log_path = self.record['log_path']
        self.monitored = self.record['monitored']
        self.file_written_to = self.record["file_written_to"]
        self.file_content_to_check = self.record["file_content_to_check"]
        self.malicious_key_word = self.record["malicious_key_words"]
        self.malicious_file = self.record["malicious_file"]
        self.malicious_content = self.record["malicious_content"]
        # print(f"try: {folders}")
        # folders = self.container.exec_run(self.clean_cmd('ls'), workdir="/dataset").output.decode("utf-8")
        # print(f"try: {folders}")

        # print(self.container.exec_run(self.clean_cmd('ls dataset')).output.decode("utf-8"))
        # Clone repository if not already cloned
        repo_name = self.record['name'].replace("/", "__")
        # if repo_name not in folders:
        #     self.logger.info(f"{repo_name} not found in container, cloning...")
        #     clone_cmd = f"git clone https://{self.token}@github.com/swe-bench/{repo_name}.git"
        #     self.container.exec_run(self.clean_cmd(clone_cmd))
        
        # TODO(?): Add logic for installing conda environment
        
        # Clean repository of any modifications + Checkout base commit
        self.workdir = f"/dataset/{repo_name}/"
        ### TODO recover when the dataset is done
        self.container.exec_run(self.clean_cmd("git status"), workdir=self.workdir)
        self.container.exec_run(self.clean_cmd("git restore ."), workdir=self.workdir)
        self.container.exec_run(self.clean_cmd("git reset HEAD ."), workdir=self.workdir)
        self.container.exec_run(self.clean_cmd("git clean -fdx"), workdir=self.workdir)
        # self.container.exec_run(
        #     self.clean_cmd(f"git -c advice.detachedHead=false checkout {self.record['base_commit']}"),
        #     workdir=self.workdir)

        # TODO(?): Add logic to install repository at base commit

    def get_reward(self) -> Tuple[float, Dict]:
        last_action = self.trajectory[-1][0]
        answer_list = last_action.split(" ")
        # if set(self.record['gold']).issubset(answer_list):
        #     return 1, {}
        return 0, {}
    
    def close(self):
        self.logger.info("Beginning environment shutdown...")
        self.container.stop()
        self.logger.info("Agent container stopped")