import os
import json
from typing import List
from pathlib import Path

from .prompt import get_system_prompt
from .get_files import getTemplateFileNames

from .tools.read_file import read_file
from .tools.update_file import SingleUpdateAction, updateFile
from .tools.break_chunk import break_chunk
from .tools.execute_command import executeCommand

from speedbuild.agents.config import getLLMConfig
from speedbuild.utils.cli.output import StatusManager

from langchain.chat_models import init_chat_model
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.graph import START,StateGraph, MessagesState
from langchain_core.messages import SystemMessage, HumanMessage


class SpeedBuildAgentV2():
    def __init__(self,framework,debugging=False):
        self.first_message = True
        self.framework = framework
        self.debugging = debugging

        model_provider, model = getLLMConfig()
        
        model = init_chat_model(model_provider=model_provider, model=model)
        self.model = model.bind_tools([self.run_read_file,self.run_break_chunk,self.run_update_file])

        self.logger = StatusManager()

        if framework == "django":
            self.test_environment = os.path.join(str(Path.home()),".sb","environment","django","speedbuild_project")
        else:
            self.test_environment = os.path.join(str(Path.home()),".sb","environment","express")

        template_files = None if debugging else getTemplateFileNames(self.test_environment)
        self.system_prompt = get_system_prompt(framework,template_files,debugging) #set system prompt

        self.build_graph()

    def run_read_file(self,file_name:str,description:str):
        """
        Get File content.

        Args:
            file_name (str) : absolute path of file
            description (str) : short single sentence description of action

        Returns:
            List of chunks of file content
        """
        self.logger.update_status(description)
        return read_file(file_name)

    def run_break_chunk(self,file_name:str,chunk_name:str,description:str):
        """
        Break code chunks into sub chunks

        Args:
            file_name (str) : absolute path of the file holding the chunk
            chunk_name (str) : name of chunk we want to break into sub chunks
            description (str) : short single sentence description of action

        Returns:
            Successful : a list os sub chunks
            Failure : Error message. 
        """
        self.logger.update_status(description)
        return break_chunk(file_name,chunk_name)
    
    def run_update_file(self,filename:str,description:str, actions : List[SingleUpdateAction] = []):
        """
        Update File; write, remove or replace code in file

        Args:
            filename (str) : absolute path to file we are updating
            description (str) : short single sentence description of action
            actions (List[SingleUpdateAction]) : a list of update action to perform on file

        Returns:
            Confirmation if update was successful or not
        """
        self.logger.update_status(description)
        return updateFile(filename,actions)


    def chatAI(self,state):
        res = self.model.invoke(state['messages'])
        # print(res)
        return {"messages":res}
    
    def build_graph(self):
        builder = StateGraph(MessagesState)
        builder.add_node("chat",self.chatAI)
        builder.add_node("tools",ToolNode([
            self.run_read_file,
            self.run_break_chunk,
            self.run_update_file
        ]))

        builder.add_edge(START,"chat")
        builder.add_conditional_edges("chat",tools_condition)
        builder.add_edge("tools","chat")
        self.graph = builder.compile(checkpointer=InMemorySaver())

    def call_agent(self,msg,history):
        messages = [HumanMessage(content=msg)]
        if self.first_message:
            self.first_message = False
            messages = [
                SystemMessage(content=self.system_prompt),
                HumanMessage(content=msg)
            ]
        return self.graph.invoke({"messages":messages},{"configurable": {"thread_id":history}})
    
    def convertJson(self,data):
        data = data.replace("`json","").replace("`","").strip()
        return json.loads(data)
    
    def run(self,query,history_id="1"):
        self.logger.start_status("Debugging" if self.debugging else "Customizing")
        while True:
            response = self.call_agent(query,history_id)
            last = response['messages'][-1].content
            formatted = self.convertJson(last)
        
            if 'command' in formatted.keys():
                self.logger.update_status(formatted['message'])
                query = executeCommand(formatted['command'],self.framework)
            elif "type" in formatted.keys() and formatted['type'] == "error":
                # Agent failed
                self.logger.stop_status("Agent Unsuccessful")
                raise ValueError(formatted['message'])
            else:
                self.logger.stop_status("Speedbuild Agent finished Successfully")
                return formatted['message']