# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.


from openenv.core.env_server.types import Action, Observation, State
from pydantic import Field,BaseModel,ConfigDict
from typing import Optional,Literal,Any
import pandas as pd
import os, yaml

class init_Directory(dict):
    def __getattr__(self, name):
        if name in self: return self[name]
        raise AttributeError(name)
    def __setattr__(self, name, value):
        self[name] = value

    def __init__(self,root,dir_list):
        super().__init__()
        setattr(self,root,self.init_dir([],dir_list))

    def init_dir(self,l,dir_list):
        if isinstance(dir_list,list):
            for i in dir_list:
                l=self.init_dir(l,i)
        if isinstance(dir_list,dict):
            for k,v in dir_list.items():
                self.create_dir(k,v)
                l.append(k+"/")
        if isinstance(dir_list,str):
            try:
                file_path = None
                possible_filename = dir_list[::-1].replace("_", ".", 1)[::-1]
                
                # Get the base directory dynamically to make it CWD independent
                base_dir = os.path.dirname(os.path.abspath(__file__))
                
                # Search for the file in the directory tree starting from the script's directory
                for root_path, _, files in os.walk(base_dir):
                    if dir_list in files:
                        file_path = os.path.join(root_path, dir_list)
                        break
                    elif possible_filename in files:
                        file_path = os.path.join(root_path, possible_filename)
                        break
                
                if not file_path:
                    file_path = dir_list

                # Read according to file type
                if file_path.endswith(('.yaml', '.yml')) or dir_list.endswith(('_yaml', '_yml')):
                    with open(file_path, "r",encoding="utf-8") as f:
                        contents = f.read()
                elif file_path.endswith('.csv') or dir_list.endswith('_csv'):
                    contents = (pd.read_csv(file_path,encoding="utf-8")).to_string(index=False)
                else:
                    with open(file_path, "r",encoding="utf-8") as f:
                        contents = f.read()
                        
                setattr(self, dir_list, contents)
                l.append(dir_list)
            except Exception:
                pass
        return l
    def create_dir(self,root,dir_list):
        setattr(self,root,self.init_dir([],dir_list))

class MyAction(Action):
    action_type: Literal[
        "cd",
        "read",
        "write",
        "run",
        "finish"
    ]= Field(
        description='''Following is a discription of the available action types:
        - "cd":Takes a path as string to change the current directory to the specified path.
        - "read":Takes a file path as string to read the content of the file at the specified path.
        - "write":Takes a file path and content as strings to write the content to the file at the specified path. If the file does not exist, it will be created, if it exists, it will be overwritten.
        - "run":Takes the path of the config file as string and runs the main python training script using that config file.
        - "finish":Takes no additional parameters and indicates that the agent has completed its task and wants to finish the episode.
        '''
    )
    path: Optional[str] = Field(default=None, description="The path to change to, read from, or write to.")
    content: Optional[str] = Field(default=None, description="The content to write or output.")


class MyObservation(Observation):
    task_type:Optional[str] = Field(description="The type of the task from easy, medium, hard.")
    task_description:Optional[str] = Field(description="A detailed description of the task to be completed.")
    cwd:Optional[str] = Field(description="The current working directory.")
    cwd_contents:Optional[list[str]] = Field(description="A list of files and directories in the current working directory.")
    file_contents:Optional[str] = Field(default=None, description="The content of the file requested by read action_type.")
    message:Optional[str] = Field(default=None, description="A message providing feedback on the last action taken, such as success or error messages.")
    remaining_steps:Optional[int] = Field(description="The number of steps remaining before the episode ends.")
    available_actions:Optional[list[str]] = Field(description="A list of available action types that the agent can take in the current state.")
    # done
    # reward
class MyState(State):
    cwd:Optional[str] = Field(description="The current working directory.")
    max_steps:Optional[int] = Field(description="The maximum number of steps allowed in the episode.")
    task_type:Optional[str] = Field(description="The type of the task from easy, medium, hard.")
    all_task_details:Optional[dict[str, list[dict[str, list[bool | float]] | str]]] = Field(description="The checklist for the task to be completed.")
    dir_structure: Any = Field(description="The structure of the directories and files in the environment.")
    cwd_contents:Optional[list[str]] = Field(description="A list of files and directories in the current working directory.")
    file_contents:Optional[str] = Field(default=None, description="The content of the file requested by read action_type.")
    # episode_id
    # step_count



    #     l=[]
    #     for dir_name in dir_list:
    #         if dir_name[-1] =='/':
    #             self.create_dir(dir_name)
    #         else:
    #             setattr(self,dir_name,"")
    #         l.append(getattr(self,dir_name))
                
    #     return l
    # def create_dir(self,dir_name):
    #     setattr(self,dir_name,[])
# class File():
#     def __init__(self,file_name,contents):
#         setattr(self,file_name,contents)