import os
from typing import Dict
from primisai.nexus.core import AI
from primisai.nexus.architect.schemas import WorkflowDefinition


class WorkflowStructurer:
    """
    Handles the structuring phase of the "Archtect" agentic workflow creation.

    This class is responsible for taking a high-level, natural language user
    prompt (an "expanded prompt") and converting it into a formal, structured
    JSON output. It achieves this by using a Large Language Model (LLM) that is
    guided by a 'NEXUS_DOCUMENTATION.md' file. This file defines the available
    components and rules for building a valid workflow.

    The structured output generated by this class is designed to be consumed by
    a "Builder" component, which then translates the structure into executable code.
    """

    def __init__(self, llm_config: Dict[str, str]):
        """
        Initializes the WorkflowStructurer instance.

        This constructor configures the internal AI client with the provided
        LLM settings. It also locates and loads the critical 'NEXUS_DOCUMENTATION.md'
        file from the same directory. This documentation is stored in memory
        to be used as context for the LLM during the structuring process.

        Args:
            llm_config (Dict[str, str]): A dictionary containing the configuration
                for the Language Model client, which may include API keys, model
                names, endpoints, etc.

        Raises:
            FileNotFoundError: If the 'NEXUS_DOCUMENTATION.md' file is not found
                in the same directory as this module.
        """
        self.llm_config = llm_config
        self.ai = AI(self.llm_config)

        curr_dir = os.path.dirname(os.path.abspath(__file__))
        doc_path = os.path.join(curr_dir, "NEXUS_DOCUMENTATION.md")

        with open(doc_path, "r", encoding="utf-8") as file:
            self.nexus_documentation = file.read()

    def reasoning_workflow_design(self, expanded_workflow: str) -> WorkflowDefinition:
        """
        Convert expanded workflow description into structured component definitions.
        
        Args:
            expanded_workflow: Detailed workflow description from previous step
            
        Returns:
            WorkflowDefinition: Structured workflow components
        """
        messages = [{
            "role":
                "system",
            "content": ("You are a workflow structure expert. Your task is to convert "
                        "the expanded workflow description into structured component "
                        "definitions following the provided schema. Include validation "
                        "constraints for each component. Don't Use Sub Supervisors.")
        }, {
            "role":
                "system",
            "content": (f"Nexus Documentation:\n\n {self.nexus_documentation}"
                        "Based on the examples provided in the Nexus documentation, "
                        "please structure the expanded workflow description into "
                        "component definitions. Include supervisors, agents, and tools. "
                        "Also, add validation constraints for each component.")
        }, {
            "role":
                "user",
            "content": (
                f"Given the following workflow description:\n\n{expanded_workflow}\n\n"
                "Please provide structured definitions for all components including "
                "supervisors, agents, and tools. For tools, include both metadata "
                "and complete Python function implementations. Also include validation "
                "constraints for each component. Also try not to add whitespace in the "
                "names of components. Add a proper (detailed) system messgages for all the components. Detailing all the details. Covering all the points. USE same system messages provided in query. DONT change them"
            )
        }]

        try:
            completion = self.ai.client.beta.chat.completions.parse(messages=messages,
                                                                    response_format=WorkflowDefinition,
                                                                    model=self.llm_config["model"])
            return completion.choices[0].message.parsed
        except Exception as e:
            raise Exception(f"Error in workflow structuring: {str(e)}")
