from typing import List,Any,Dict
from openai import OpenAI
from typing import Optional
import json
import re
from HeGFlow.llm.llm_registry import LLMRegistry
from HeGFlow.graph.node import Node
from HeGFlow.agents.agent_registry import AgentRegistry
from HeGFlow.llm.llm_registry import LLMRegistry
from HeGFlow.prompt.prompt_set_registry import PromptSetRegistry
from HeGFlow.tools.coding.python_executor import execute_code_get_return
from local_datasets.gsm8k_dataset import gsm_get_predict
from HeGFlow.tools.search.Serper import search_serper_main
from HeGFlow.tools.General.general_tools import str_to_list

@AgentRegistry.register('CriticNode')
class CriticNode(Node):
    """
    A special node that evaluates outputs of other nodes by role, generating feedback.
    Does not participate in spatial/temporal predecessor/successor connections.
    """
    def __init__(self, id: Optional[str] = None, agent_name: str = "CriticNode", domain: str = "", llm_name: str = ""):
        super().__init__(id, agent_name, domain, llm_name)
        self.critic_connections: List[Dict[str, Any]] = []  
        self.role = "CriticNode"
        self.llm = LLMRegistry.get(llm_name)

    def add_critic_connection(self, node: Node):  
        """Add a node to critic_connections without affecting spatial/temporal connections."""
        if not any(conn["node_id"] == node.id for conn in self.critic_connections):
            self.critic_connections.append({"node_id": node.id, "node": node})
            node.critic_enabled  = True
    def printtxt(self):
        pass

    def remove_critic_connection(self, node: Node):  
        """Remove a node from critic_connections."""
        self.critic_connections = [conn for conn in self.critic_connections if conn["node_id"] != node.id]
        node.critic_enabled = False

    def extract_json(self,text: str) -> str:
       
        json_match = re.search(r"```(?:json)?\n(.*)```", text, re.DOTALL)
        if json_match:
            return json_match.group(1).strip()  
        return text 

    def _execute(self, input: List[Any], spatial_info: Dict[str, Any], temporal_info: Dict[str, Any], **kwargs):  
        raise NotImplementedError("Synchronous execution not supported for CriticNode")

    async def _async_execute(self, input: List[Any], spatial_info: Dict[str, Any], temporal_info: Dict[str, Any], **kwargs):
        """
        Evaluate outputs of connected nodes by role.
        Input: {task: str, role_outputs: {role: {node_id: output}}}
        Output: [{node_id: str, feedback: {issues: [], suggestions: []}}]
        """
        task = input["task"] 
        role_outputs = input.get("role_outputs", {})
        feedback = []  
        for role, outputs in role_outputs.items():
            for node_id, output in outputs.items():
               
                node = next((conn["node"] for conn in self.critic_connections if conn["node_id"] == node_id), None)
                if not node or not node.critic_enabled:
                    if not node:
                        print(f"Node with id {node_id} not found.")
                    else:
                        print(f"Node with id {node_id} has critic_enabled=False.")
                    print("Skipping node with critic_enabled=False or no output.")
                    continue
                if output == "None." or not output:
                    print("Skipping node with empty output.")
                    continue
               
                prompt = f"""
                Task: {task}
                Role: {role}
                Output: {output}
                ## Evaluation Criteria:

                You are an expert evaluator. Your goal is to critically assess the provided output based on the following criteria, tailored to the specific task and role. Consider edge cases, potential ambiguities, and subtle errors.

                1.   **Correctness:** Is the information factually accurate and free from errors? (Especially important for factual tasks)
                2.   **Completeness:** Does the output fully address the task requirements? Are there any missing pieces of information or steps?
                3.   **Logical Consistency:** Is the reasoning sound and coherent? Are there any contradictions or logical fallacies?  Does the solution make sense in the context of the task?
                4.   **Relevance:** Is the output focused on the task and role, or does it contain irrelevant or extraneous information?
                5.   **Clarity:** Is the output easy to understand and well-organized? Is the language precise and unambiguous?

                ## Evaluation Process:

                1.  **Deep Dive:** Carefully analyze the task, role, and output.
                2.  **Error Detection:** Identify any specific issues or areas where the output falls short based on the above criteria.
                3.  **Suggestion Generation:** Provide concrete, actionable suggestions for improving the output. Be specific about *what* needs to be changed and *why*.
                
                Return feedback as JSON: {{"issues": [], "suggestions": []}}.
                when you believe this answer is correct, simply set the "issues" field in the output to an empty list.
                Example:
                - Input: "Solve 2 + 2", Output: "5"
                  Feedback: {{"issues": ["Incorrect calculation: 2 + 2 = 4, not 5"], "suggestions": ["Recalculate the sum"]}}
                -  Input: "Solve 4 * 6", Output: "According to the multiplication calculation, 4 * 6 = 24"
                  Feedback: {{"issues": [], "suggestions": []}}
                -  Input: "What year did World War II begin?", Output: "World War II started in 1941."
                  Feedback: {{"issues": ["Incorrect start date: WWII began in 1939, not 1941."], "suggestions": ["Verify the start date from a reliable source and correct it."]}}
                -  Input: "Why does the Earth have seasons?", Output: "Because the Earth orbits the sun."
                  Feedback: {{"issues": ["Incomplete explanation: The Earth's tilt is the primary cause."], "suggestions": ["Mention the Earth's axial tilt as the main reason for the seasons."]}}
                In order for the model to produce more accurate output, please give feedback as frequently as possible.Only when the model's reasoning and answers are impeccable, give empty feedback
                """
                
                response = await self._call_llm(prompt)  # Placeholder for LLM call
                response = self.extract_json(response)

                feedback.append({"node_id": node_id, "feedback": json.loads(response)})
        #self.outputs = feedback
        return feedback
    async def _async_execute_single_node(self, input: List[Any], spatial_info: Dict[str, Any], temporal_info: Dict[str, Any], **kwargs):
        task = input["task"]
        role_outputs = input.get("role_outputs", {})
        role = next(iter(role_outputs))
        outputs = role_outputs[role]

      
        node_id = next(iter(outputs)) 
        output_text = outputs[node_id]
        output = output_text
        feedback = [] 
        prompt = f"""
                Task: {task}
                Role: {role}
                Output: {output}
                ## Evaluation Criteria:

                You are an expert evaluator. Your goal is to critically assess the provided output based on the following criteria, tailored to the specific task and role. Consider edge cases, potential ambiguities, and subtle errors.

                1.   **Correctness:** Is the information factually accurate and free from errors? (Especially important for factual tasks)
                2.   **Completeness:** Does the output fully address the task requirements? Are there any missing pieces of information or steps?
                3.   **Logical Consistency:** Is the reasoning sound and coherent? Are there any contradictions or logical fallacies?  Does the solution make sense in the context of the task?
                4.   **Relevance:** Is the output focused on the task and role, or does it contain irrelevant or extraneous information?
                5.   **Clarity:** Is the output easy to understand and well-organized? Is the language precise and unambiguous?

                ## Evaluation Process:

                1.  **Deep Dive:** Carefully analyze the task, role, and output.
                2.  **Error Detection:** Identify any specific issues or areas where the output falls short based on the above criteria.
                3.  **Suggestion Generation:** Provide concrete, actionable suggestions for improving the output. Be specific about *what* needs to be changed and *why*.
                
                Return feedback as JSON: {{"issues": [], "suggestions": []}}.
                when you believe this answer is correct, simply set the "issues" field in the output to an empty list.
                Example:
                - Input: "Solve 2 + 2", Output: "5"
                  Feedback: {{"issues": ["Incorrect calculation: 2 + 2 = 4, not 5"], "suggestions": ["Recalculate the sum"]}}
                -  Input: "Solve 4 * 6", Output: "According to the multiplication calculation, 4 * 6 = 24"
                  Feedback: {{"issues": [], "suggestions": []}}
                -  Input: "What year did World War II begin?", Output: "World War II started in 1941."
                  Feedback: {{"issues": ["Incorrect start date: WWII began in 1939, not 1941."], "suggestions": ["Verify the start date from a reliable source and correct it."]}}
                -  Input: "Why does the Earth have seasons?", Output: "Because the Earth orbits the sun."
                  Feedback: {{"issues": ["Incomplete explanation: The Earth's tilt is the primary cause."], "suggestions": ["Mention the Earth's axial tilt as the main reason for the seasons."]}}
                In order for the model to produce more accurate output, please give feedback as frequently as possible.Only when the model's reasoning and answers are impeccable, give empty feedback
                """
        
        response = await self._call_llm(prompt)  # Placeholder for LLM call
        response = self.extract_json(response)
        feedback.append({"node_id": node_id, "feedback": json.loads(response)})
        return feedback
    

    async def _call_llm(self, prompt: str) -> str:
        """
        Placeholder for async LLM call. Replace with actual LLM interface.
        """
        # Simulate LLM response (replace with actual implementation)
        client = OpenAI(api_key="", base_url="")
        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                    {"role": "system", "content": prompt},
                    {"role": "user", "content": ""},
                ],
            stream=False
        )
        response_str = response.choices[0].message.content
        return response_str
        

    def _process_inputs(self, raw_inputs: List[Any], spatial_info: Dict[str, Any], temporal_info: Dict[str, Any], **kwargs) -> List[Any]:
        """
        Process inputs to extract task and role_outputs.
        """
        if not raw_inputs or not isinstance(raw_inputs, dict):
            return []
        return [raw_inputs]