// Copyright (c) Microsoft. All rights reserved.

using AgentWebChat.AgentHost;
using AgentWebChat.AgentHost.Custom;
using AgentWebChat.AgentHost.Utilities;
using Microsoft.Agents.AI;
using Microsoft.Agents.AI.DevUI;
using Microsoft.Agents.AI.Hosting;
using Microsoft.Agents.AI.Workflows;
using Microsoft.Extensions.AI;

var builder = WebApplication.CreateBuilder(args);

// Add service defaults & Aspire client integrations.
builder.AddServiceDefaults();
builder.Services.AddOpenApi();

// Add services to the container.
builder.Services.AddProblemDetails();

// Configure the chat model and our agent.
builder.AddKeyedChatClient("chat-model");

// Add DevUI services
builder.AddDevUI();

// Add OpenAI services
builder.AddOpenAIChatCompletions();
builder.AddOpenAIResponses();

// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions by authenticated caller.
// Without this, contextId alone is the session key — any caller who knows a contextId can access that session.
// Example using claims-based identity:
// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });

// By default, NoopAgentSessionStore is used — sessions are not persisted across requests.
// To enable multi-turn conversations, register a session store explicitly, e.g.:
// agentBuilder.WithInMemorySessionStore();

var pirateAgentBuilder = builder.AddAIAgent(
    "pirate",
    instructions: "You are a pirate. Speak like a pirate",
    description: "An agent that speaks like a pirate.",
    chatClientServiceKey: "chat-model")
    .WithAITool(new CustomAITool())
    .WithAITool(new CustomFunctionTool())
    .WithInMemorySessionStore();

var knightsKnavesAgentBuilder = builder.AddAIAgent("knights-and-knaves", (sp, key) =>
{
    var chatClient = sp.GetRequiredKeyedService<IChatClient>("chat-model");

    ChatClientAgent knight = new(
        chatClient,
        """
        You are a knight. This means that you must always tell the truth. Your name is Alice.
        Bob is standing next to you. Bob is a knave, which means he always lies.
        When replying, always start with your name (Alice). Eg, "Alice: I am a knight."
        """, "Alice");

    ChatClientAgent knave = new(
        chatClient,
        """
        You are a knave. This means that you must always lie. Your name is Bob.
        Alice is standing next to you. Alice is a knight, which means she always tells the truth.
        When replying, always include your name (Bob). Eg, "Bob: I am a knight."
        """, "Bob");

    ChatClientAgent narrator = new(
        chatClient,
        """
        You are are the narrator of a puzzle involving knights (who always tell the truth) and knaves (who always lie).
        The user is going to ask questions and guess whether Alice or Bob is the knight or knave.
        Alice is standing to one side of you. Alice is a knight, which means she always tells the truth.
        Bob is standing to the other side of you. Bob is a knave, which means he always lies.
        When replying, always include your name (Narrator).
        Once the user has deduced what type (knight or knave) both Alice and Bob are, tell them whether they are right or wrong.
        If the user asks a general question about their surrounding, make something up which is consistent with the scenario.
        """, "Narrator");

    return AgentWorkflowBuilder.BuildConcurrent([knight, knave, narrator]).AsAIAgent(name: key);
});

// Workflow consisting of multiple specialized agents
var chemistryAgent = builder.AddAIAgent("chemist",
    instructions: "You are a chemistry expert. Answer thinking from the chemistry perspective",
    description: "An agent that helps with chemistry.",
    chatClientServiceKey: "chat-model");

var mathsAgent = builder.AddAIAgent("mathematician",
    instructions: "You are a mathematics expert. Answer thinking from the maths perspective",
    description: "An agent that helps with mathematics.",
    chatClientServiceKey: "chat-model");

var literatureAgent = builder.AddAIAgent("literator",
    instructions: "You are a literature expert. Answer thinking from the literature perspective",
    description: "An agent that helps with literature.",
    chatClientServiceKey: "chat-model");

var scienceSequentialWorkflow = builder.AddWorkflow("science-sequential-workflow", (sp, key) =>
{
    List<IHostedAgentBuilder> usedAgents = [chemistryAgent, mathsAgent, literatureAgent];
    var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService<AIAgent>(ab.Name));
    return AgentWorkflowBuilder.BuildSequential(workflowName: key, agents: agents);
}).AddAsAIAgent();

var scienceConcurrentWorkflow = builder.AddWorkflow("science-concurrent-workflow", (sp, key) =>
{
    List<IHostedAgentBuilder> usedAgents = [chemistryAgent, mathsAgent, literatureAgent];
    var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService<AIAgent>(ab.Name));
    return AgentWorkflowBuilder.BuildConcurrent(workflowName: key, agents: agents);
}).AddAsAIAgent();

builder.AddWorkflow("nonAgentWorkflow", (sp, key) =>
{
    List<IHostedAgentBuilder> usedAgents = [pirateAgentBuilder, chemistryAgent];
    var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService<AIAgent>(ab.Name));
    return AgentWorkflowBuilder.BuildSequential(workflowName: key, agents: agents);
});

builder.Services.AddKeyedSingleton("NonAgentAndNonmatchingDINameWorkflow", (sp, key) =>
{
    List<IHostedAgentBuilder> usedAgents = [pirateAgentBuilder, chemistryAgent];
    var agents = usedAgents.Select(ab => sp.GetRequiredKeyedService<AIAgent>(ab.Name));
    return AgentWorkflowBuilder.BuildSequential(workflowName: "random-name", agents: agents);
});

builder.Services.AddSingleton<AIAgent>(sp =>
{
    var chatClient = sp.GetRequiredKeyedService<IChatClient>("chat-model");
    return new ChatClientAgent(chatClient, name: "default-agent", instructions: "you are a default agent.");
});

builder.Services.AddKeyedSingleton<AIAgent>("my-di-nonmatching-agent", (sp, name) =>
{
    var chatClient = sp.GetRequiredKeyedService<IChatClient>("chat-model");
    return new ChatClientAgent(
        chatClient,
        name: "some-random-name", // demonstrating registration can be different for DI and actual agent
        instructions: "you are a dependency inject agent. Tell me all about dependency injection.");
});

builder.Services.AddKeyedSingleton<AIAgent>("my-di-matchingname-agent", (sp, name) =>
{
    if (name is not string nameStr)
    {
        throw new NotSupportedException("Name should be passed as a key");
    }

    var chatClient = sp.GetRequiredKeyedService<IChatClient>("chat-model");
    return new ChatClientAgent(
        chatClient,
        name: nameStr, // demonstrating registration with the same name
        instructions: "you are a dependency inject agent. Tell me all about dependency injection.");
});

pirateAgentBuilder.AddA2AServer();
knightsKnavesAgentBuilder.AddA2AServer();

// IMPORTANT: In production, register a SessionIsolationKeyProvider to isolate sessions by authenticated caller.
// Without this, contextId alone is the session key — any caller who knows a contextId can access that session.
// Example using claims-based identity:
// builder.Services.UseClaimsBasedSessionIsolation(new() { ClaimType = ClaimTypes.NameIdentifier });

var app = builder.Build();

app.MapOpenApi();
app.UseSwaggerUI(options => options.SwaggerEndpoint("/openapi/v1.json", "Agents API"));

// Configure the HTTP request pipeline.
app.UseExceptionHandler();

// Expose A2A servers over HTTP with JSON payloads
app.MapA2AHttpJson(pirateAgentBuilder, path: "/a2a/pirate");
app.MapA2AHttpJson(knightsKnavesAgentBuilder, path: "/a2a/knights-and-knaves");

app.MapDevUI();

app.MapOpenAIResponses();
app.MapOpenAIResponses(pirateAgentBuilder);
app.MapOpenAIResponses(knightsKnavesAgentBuilder);
app.MapOpenAIResponses(chemistryAgent);
app.MapOpenAIResponses(mathsAgent);
app.MapOpenAIResponses(literatureAgent);
app.MapOpenAIResponses(scienceSequentialWorkflow);
app.MapOpenAIResponses(scienceConcurrentWorkflow);
app.MapOpenAIConversations();

app.MapOpenAIChatCompletions(pirateAgentBuilder);
app.MapOpenAIChatCompletions(knightsKnavesAgentBuilder);
app.MapOpenAIChatCompletions(chemistryAgent);
app.MapOpenAIChatCompletions(mathsAgent);
app.MapOpenAIChatCompletions(literatureAgent);
app.MapOpenAIChatCompletions(scienceSequentialWorkflow);
app.MapOpenAIChatCompletions(scienceConcurrentWorkflow);

// Map the agents HTTP endpoints
app.MapAgentDiscovery("/agents");

app.MapDefaultEndpoints();
app.Run();
