from typing import Dict
from pydantic import BaseModel, Field
from cognify.llm import Input, Model, StructuredModel, OutputFormat
from .prompts import decompose_system


class AgentPropose(BaseModel):
    """Proposed agent metadata"""

    available_information: str = Field(
        description="information that the agent can receive"
    )
    prompt: str = Field(description="prompt for the agent")
    output: str = Field(description="content generated by the agent")


class HighLevelDecompose(BaseModel):
    """High level decomposition of the task"""

    agents: Dict[str, AgentPropose] = Field(
        description="Dictionary of agent name to agent metadata"
    )


def high_level_decompose_kernel(task: str) -> HighLevelDecompose:
    task_input = Input(name="task")
    construct_new_agents = Model(
        "decompose", decompose_system, input_variables=[task_input]
    )
    messages = [
        {"role": "user", "content": f"Original Single Agent Prompt:\n{task}\n\n"}
    ]
    model_kwargs = {"model": "gpt-4o", "temperature": 0.0}
    new_agents_raw: str = construct_new_agents(
        messages, model_kwargs, inputs={task_input: task}
    )

    # format to pydantic object
    new_agents_input = Input(name="new_agents")
    format_agent = StructuredModel(
        "decompose_struct",
        decompose_system,
        input_variables=[new_agents_input],
        output_format=OutputFormat(
            schema=HighLevelDecompose, should_hint_format_in_prompt=True
        ),
    )
    messages.extend(
        [
            {"role": "assistant", "content": f"{new_agents_raw}"},
            {
                "role": "user",
                "content": "Now please reformat your answer to the JSON format.\n",
            },
        ]
    )
    new_agents: HighLevelDecompose = format_agent(
        messages, model_kwargs, inputs={new_agents_input: new_agents_raw}
    )
    new_agents = HighLevelDecompose.model_validate(new_agents)
    return new_agents
