from pydantic_ai.models.openai import OpenAIModel
from core.azure import azure_provider  
from pydantic_ai import Agent
from pydantic_ai.settings import ModelSettings

async def StructuredOutputAgent(model_name, raw_response, response_format):
    completion = Agent(
        model=OpenAIModel(model_name, provider=azure_provider),
        output_type=response_format
    )

    task_description = f"given the content: {raw_response}, rewrite exactly the structured output writen in content"

    result = await completion.run(task_description, model_settings=ModelSettings(temperature=0.0, seed=100))
    if result is None:
        raise Exception("No result from get_agent_graph_entities")
    return result.output

async def run_agent(model_name, system_prompt, task_description, response_format):
    agent = Agent(
        model=OpenAIModel(model_name, provider=azure_provider),
        system_prompt=system_prompt,
    )
    result = await agent.run(task_description, model_settings=ModelSettings(temperature=0.5, seed=100))
    raw_response = result.output

    if response_format:
        return await StructuredOutputAgent(model_name, raw_response, response_format)
    else:
        return raw_response


def create_system_prompt(role, backstory, goal):
    return f"""You are {role}. {backstory} \n Your goal is {goal}"""


def create_task_prompt(task_description, task_expected_output, response_format, previous_step_output=None):
    return f"""Let's think step by step to ensure logical and evidence-based reasoning.

## Task: {task_description}

## Context from previous task: {previous_step_output if previous_step_output else "None provided."}

Follow this step-by-step plan of action:
1. Understand the problem: Summarize the main context and identify key elements.
2. Break it down: Divide the task into logical sub-steps or components.
3. Gather information: List relevant facts, data, or assumptions (if applicable, search or analyze inputs).
4. Generate options: Explore multiple ideas or paths (e.g., 3-5 alternatives, considering branching paths like in Tree-of-Thoughts).
5. Evaluate: Analyze pros, cons, and feasibility of each option, using self-consistency (generate 2-3 variants if uncertain).
6. Decide and act: Choose the best option and detail an action plan with timelines or responsibilities.
7. Reflect: Evaluate confidence in the response (High/Medium/Low) and explain why; suggest improvements if needed.

Output in this structured format:
- **Step 1: [Step Title]**  
  [Detailed content].
- **Step 2: [Step Title]**  
  [Detailed content].
...
- **Final Plan of Action:**  
  [Summary of the plan, e.g., in bullets or numbered list].
- **Reflection:**  
  [Evaluation].

Keep the response concise, under 500 tokens. Use evidence-based reasoning and cite sources if available.

Expected output: {task_expected_output}

Finally, after thinking step by step, answer in this format: {response_format}
"""