# echo_agent.py
import os
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Microsoft Agent Framework Imports (Conceptual/Based on .NET structure)
# Note: The Python SDK is still evolving. This uses currently available modules
# or conceptually aligns with the .NET MAF structure and general Python SDK patterns.
from microsoft.agents.hosting.core import AgentApplication, TurnState, TurnContext, MemoryStorage
from microsoft.agents.hosting.aiohttp import CloudAdapter, start_agent_process, jwt_authorization_middleware
from aiohttp.web import Request, Response, Application, run_app

# OpenAI / Azure OpenAI imports
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential

# --- Configuration for Azure OpenAI ---
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME")

if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_DEPLOYMENT_NAME]):
    raise ValueError("Missing Azure OpenAI environment variables. Please check your .env file.")

# Authenticate with Azure OpenAI
# For production, consider DefaultAzureCredential for managed identities
try:
    aoai_client = AzureOpenAI(
        api_key=AZURE_OPENAI_API_KEY,
        api_version="2024-02-15-preview", # Check latest API version
        azure_endpoint=AZURE_OPENAI_ENDPOINT,
    )
    print("Successfully connected to Azure OpenAI.")
except Exception as e:
    print(f"Error connecting to Azure OpenAI: {e}")
    # Fallback or exit if essential AI service is unavailable
    exit(1)

# A simple wrapper for an LLM client that conforms to MAF agent expectations
class AzureOpenAIChatClient:
    def __init__(self, client, deployment_name):
        self.client = client
        self.deployment_name = deployment_name

    async def complete(self, messages: list[dict], **kwargs) -> str:
        try:
            response = await self.client.chat.completions.create(
                model=self.deployment_name,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error calling Azure OpenAI: {e}")
            return "An error occurred while processing your request."

# Initialize our custom chat client
my_chat_client = AzureOpenAIChatClient(aoai_client, AZURE_OPENAI_DEPLOYMENT_NAME)

# --- Define our Agent using MAF Python SDK (Echo Agent) ---
# The AgentApplication is the core entity that handles incoming activities
# and routes them to appropriate handlers (message, conversation_update, etc.)

# In MAF Python SDK, agent logic is often implemented as handlers for TurnContext activities.
# The TurnState can store conversational state.
class EchoAgentState(TurnState):
    # You can add state variables here
    pass

AGENT_APP = AgentApplication[EchoAgentState](
    storage=MemoryStorage(), # Simple in-memory storage for demonstration
    adapter=CloudAdapter()
)

# Handler for new conversations or members added
async def _welcome_and_help(context: TurnContext, _: EchoAgentState):
    await context.send_activity(
        "Hello! I am an Echo Agent powered by Azure OpenAI. "
        "Type /help for help or send me any message to hear it echoed back."
    )

AGENT_APP.conversation_update("membersAdded")(_welcome_and_help)
AGENT_APP.message("/help")(_welcome_and_help)

# Main message handler for the Echo Agent
@AGENT_APP.activity("message")
async def on_message(context: TurnContext, _):
    user_message = context.activity.text
    # Use our Azure OpenAI client to generate a more "intelligent" echo
    # In a real agent, this would involve more complex prompt engineering
    # or tool calls.
    llm_response = await my_chat_client.complete(
        messages=[
            {"role": "system", "content": "You are a helpful assistant. Echo the user's message, but add a friendly remark."},
            {"role": "user", "content": user_message}
        ]
    )
    await context.send_activity(f"You said: '{user_message}'\nAI says: {llm_response}")


# --- FastAPI-like server setup for local testing ---
# This part hosts the agent as an HTTP endpoint, typically for local testing
# or integration with a chatbot frontend.
def start_server(agent_application: AgentApplication):
    """
    Starts an aiohttp web server to host the AgentApplication.
    """
    async def entry_point(req: Request) -> Response:
        agent: AgentApplication = req.app["agent_app"]
        adapter: CloudAdapter = req.app["adapter"]
        return await start_agent_process(req, agent, adapter)

    APP = Application() # middlewares=[jwt_authorization_middleware] for auth
    APP.router.add_post("/api/messages", entry_point)
    APP.router.add_get("/api/messages", lambda _: Response(status=200)) # Health check
    APP["agent_app"] = agent_application
    APP["adapter"] = agent_application.adapter
    try:
        port = int(os.environ.get("PORT", 3978))
        print(f"======== Running on http://localhost:{port} ========")
        print("(Press CTRL+C to quit)")
        run_app(APP, host="localhost", port=port)
    except Exception as error:
        print(f"Error starting server: {error}")
        raise error

if __name__ == "__main__":
    start_server(AGENT_APP)

# tools.py
import asyncio
from typing import Dict, Any, List

class SearchTool:
    @staticmethod
    async def web_search(query: str) -> List[str]:
        """
        Performs a simulated web search and returns relevant snippets.
        In a real application, this would integrate with a search API (e.g., Bing, Google).
        """
        print(f"Executing web search for: {query}")
        # Simulate an API call
        await asyncio.sleep(1)
        if "weather" in query.lower():
            return ["Current weather in London: 15°C, cloudy. Expected rain tomorrow."]
        elif "capital of france" in query.lower():
            return ["The capital of France is Paris."]
        elif "python agent framework" in query.lower():
            return [
                "Microsoft Agent Framework (MAF) is an open-source SDK for building AI agents.",
                "LangChain and AutoGen are other popular Python agent frameworks.",
                "MAF combines Semantic Kernel and AutoGen capabilities."
            ]
        else:
            return [f"Search results for '{query}': No specific results found, but here's some general info."]

    @staticmethod
    async def get_current_time() -> str:
        """
        Retrieves the current date and time.
        """
        import datetime
        print("Executing get_current_time tool.")
        return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# echo_agent_with_tools.py (Modified from echo_agent.py)
import os
import asyncio
from dotenv import load_dotenv

# Load environment variables
load_dotenv()

from microsoft.agents.hosting.core import AgentApplication, TurnState, TurnContext, MemoryStorage
from microsoft.agents.hosting.aiohttp import CloudAdapter, start_agent_process
from aiohttp.web import Request, Response, Application, run_app
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential

# Import our custom tools
from tools import SearchTool

# --- Configuration (same as before) ---
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME")

if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_DEPLOYMENT_NAME]):
    raise ValueError("Missing Azure OpenAI environment variables. Please check your .env file.")

try:
    aoai_client = AzureOpenAI(
        api_key=AZURE_OPENAI_API_KEY,
        api_version="2024-02-15-preview", # Check latest API version
        azure_endpoint=AZURE_OPENAI_ENDPOINT,
    )
    print("Successfully connected to Azure OpenAI.")
except Exception as e:
    print(f"Error connecting to Azure OpenAI: {e}")
    exit(1)

class AzureOpenAIChatClient:
    def __init__(self, client, deployment_name):
        self.client = client
        self.deployment_name = deployment_name

    async def complete(self, messages: list[dict], **kwargs) -> str:
        try:
            response = await self.client.chat.completions.create(
                model=self.deployment_name,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error calling Azure OpenAI: {e}")
            return "An error occurred while processing your request."

my_chat_client = AzureOpenAIChatClient(aoai_client, AZURE_OPENAI_DEPLOYMENT_NAME)

# --- Introduce Planner for Tool Calling ---
# In the MAF Python SDK, explicit tool registration and an LLM-driven planner
# are used to enable tool calling. This is where AutoGen's influence is clear.
# For simplicity, we'll create a basic planner that just decides if a search is needed.
# In a full MAF implementation, this would be handled more abstractly by the framework.

# We'll need a mechanism to describe our tools to the LLM so it knows when to use them.
# This often involves creating "function schemas" for the LLM.

def get_tool_schemas():
    """Returns a list of tool descriptions in OpenAI function call format."""
    return [
        {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "Performs a web search to find relevant information.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query."
                        }
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_current_time",
                "description": "Retrieves the current date and time.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        }
    ]

# --- Define our Agent with Tool-Calling Capability ---
class ToolUsingAgentState(TurnState):
    # We can store conversational history
    history: List[Dict[str, Any]] = []

AGENT_APP = AgentApplication[ToolUsingAgentState](
    storage=MemoryStorage(),
    adapter=CloudAdapter()
)

async def _welcome_and_help(context: TurnContext, state: ToolUsingAgentState):
    await context.send_activity(
        "Hello! I am an Agent with web search capabilities. "
        "Ask me questions like 'What is the weather in London?' or 'What is the current time?'"
    )
    state.history.append({"role": "assistant", "content": "Hello! I am an Agent with web search capabilities. Ask me questions like 'What is the weather in London?' or 'What is the current time?'"})

AGENT_APP.conversation_update("membersAdded")(_welcome_and_help)
AGENT_APP.message("/help")(_welcome_and_help)

@AGENT_APP.activity("message")
async def on_message_with_tools(context: TurnContext, state: ToolUsingAgentState):
    user_message = context.activity.text
    state.history.append({"role": "user", "content": user_message})

    messages_for_llm = [
        {"role": "system", "content": "You are a helpful assistant. Use tools if necessary to answer questions."},
        *state.history
    ]

    try:
        response = await aoai_client.chat.completions.create(
            model=AZURE_OPENAI_DEPLOYMENT_NAME,
            messages=messages_for_llm,
            tools=get_tool_schemas(),
            tool_choice="auto" # Let the LLM decide if it wants to call a tool
        )

        assistant_message = response.choices[0].message
        state.history.append(assistant_message)

        if assistant_message.tool_calls:
            for tool_call in assistant_message.tool_calls:
                function_name = tool_call.function.name
                arguments = tool_call.function.arguments

                print(f"Agent wants to call tool: {function_name} with args: {arguments}")

                # Execute the tool
                tool_output = None
                if function_name == "web_search":
                    query = eval(arguments).get("query") # Be careful with eval in production!
                    tool_output = await SearchTool.web_search(query)
                elif function_name == "get_current_time":
                    tool_output = await SearchTool.get_current_time()
                else:
                    tool_output = f"Unknown tool: {function_name}"

                state.history.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": str(tool_output)
                })

                # Call LLM again with tool output
                final_response = await aoai_client.chat.completions.create(
                    model=AZURE_OPENAI_DEPLOYMENT_NAME,
                    messages=state.history
                )
                final_assistant_message = final_response.choices[0].message.content
                state.history.append({"role": "assistant", "content": final_assistant_message})
                await context.send_activity(final_assistant_message)

        else:
            # If no tool call, just respond with the LLM's message
            await context.send_activity(assistant_message.content)

    except Exception as e:
        print(f"Error in tool-using agent: {e}")
        error_msg = "An error occurred while processing your request with tools."
        await context.send_activity(error_msg)
        state.history.append({"role": "assistant", "content": error_msg})


# --- FastAPI-like server setup (same as before) ---
def start_server(agent_application: AgentApplication):
    async def entry_point(req: Request) -> Response:
        agent: AgentApplication = req.app["agent_app"]
        adapter: CloudAdapter = req.app["adapter"]
        return await start_agent_process(req, agent, adapter)

    APP = Application()
    APP.router.add_post("/api/messages", entry_point)
    APP.router.add_get("/api/messages", lambda _: Response(status=200))
    APP["agent_app"] = agent_application
    APP["adapter"] = agent_application.adapter
    try:
        port = int(os.environ.get("PORT", 3978))
        print(f"======== Running on http://localhost:{port} ========")
        print("(Press CTRL+C to quit)")
        run_app(APP, host="localhost", port=port)
    except Exception as error:
        print(f"Error starting server: {error}")
        raise error

if __name__ == "__main__":
    start_server(AGENT_APP)

# redis_agent.py (Building on tool_agent.py)
import os
import asyncio
import json
from dotenv import load_dotenv
import redis.asyncio as redis # Use async Redis client

load_dotenv()

from microsoft.agents.hosting.core import AgentApplication, TurnState, TurnContext, MemoryStorage
from microsoft.agents.hosting.aiohttp import CloudAdapter, start_agent_process
from aiohttp.web import Request, Response, Application, run_app
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential

from tools import SearchTool # Our custom tools

# --- Configuration ---
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME")

REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")

if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_DEPLOYMENT_NAME]):
    raise ValueError("Missing Azure OpenAI environment variables. Please check your .env file.")

try:
    aoai_client = AzureOpenAI(
        api_key=AZURE_OPENAI_API_KEY,
        api_version="2024-02-15-preview",
        azure_endpoint=AZURE_OPENAI_ENDPOINT,
    )
    print("Successfully connected to Azure OpenAI.")
except Exception as e:
    print(f"Error connecting to Azure OpenAI: {e}")
    exit(1)

# Initialize Redis client
# Note: For production, use connection pooling and proper error handling.
try:
    redis_client = redis.Redis(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD,
        decode_responses=True # Decode responses to string
    )
    # Test connection
    async def _test_redis_connection():
        await redis_client.ping()
        print("Successfully connected to Redis.")
    asyncio.run(_test_redis_connection())
except Exception as e:
    print(f"Error connecting to Redis: {e}")
    # Proceed without Redis or exit based on your application's tolerance
    redis_client = None

class AzureOpenAIChatClient:
    def __init__(self, client, deployment_name):
        self.client = client
        self.deployment_name = deployment_name

    async def complete(self, messages: list[dict], **kwargs) -> str:
        try:
            response = await self.client.chat.completions.create(
                model=self.deployment_name,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error calling Azure OpenAI: {e}")
            return "An error occurred while processing your request."

my_chat_client = AzureOpenAIChatClient(aoai_client, AZURE_OPENAI_DEPLOYMENT_NAME)

def get_tool_schemas():
    """Returns a list of tool descriptions in OpenAI function call format."""
    return [
        {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "Performs a web search to find relevant information.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query."
                        }
                    },
                    "required": ["query"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_current_time",
                "description": "Retrieves the current date and time.",
                "parameters": {
                    "type": "object",
                    "properties": {},
                    "required": []
                }
            }
        }
    ]

# --- Custom Redis-backed Memory Storage (Conceptual, as MAF might have its own) ---
# For demonstration, we'll implement a simple Redis-backed storage for chat history.
# MAF's official Redis memory module would abstract this further.
class RedisChatHistoryStorage(MemoryStorage):
    def __init__(self, redis_client: redis.Redis, prefix: str = "chat_history:"):
        self.redis_client = redis_client
        self.prefix = prefix

    async def _get_key(self, conversation_id: str) -> str:
        return f"{self.prefix}{conversation_id}"

    async def read(self, conversation_id: str) -> dict:
        key = await self._get_key(conversation_id)
        raw_history = await self.redis_client.get(key)
        if raw_history:
            return json.loads(raw_history)
        return {"history": []} # Default if no history

    async def write(self, conversation_id: str, data: dict):
        key = await self._get_key(conversation_id)
        await self.redis_client.set(key, json.dumps(data))
        # Optional: Set an expiry for the history
        # await self.redis_client.expire(key, 3600) # 1 hour

# Our agent state now includes a way to load/save history
class PersistentToolUsingAgentState(TurnState):
    history: List[Dict[str, Any]] = []
    conversation_id: str = "default_conversation" # Placeholder, actual ID comes from TurnContext

# Use our Redis storage if available, otherwise fallback to in-memory
storage_instance = RedisChatHistoryStorage(redis_client) if redis_client else MemoryStorage()

AGENT_APP = AgentApplication[PersistentToolUsingAgentState](
    storage=storage_instance,
    adapter=CloudAdapter()
)

async def _load_history(context: TurnContext, state: PersistentToolUsingAgentState):
    if context.activity and context.activity.conversation and context.activity.conversation.id:
        state.conversation_id = context.activity.conversation.id
    
    if redis_client:
        stored_data = await storage_instance.read(state.conversation_id)
        state.history = stored_data.get("history", [])
        print(f"Loaded history for {state.conversation_id}: {len(state.history)} items")

async def _save_history(context: TurnContext, state: PersistentToolUsingAgentState):
    if redis_client and state.conversation_id:
        await storage_instance.write(state.conversation_id, {"history": state.history})
        print(f"Saved history for {state.conversation_id}: {len(state.history)} items")


# Register these as pre/post handlers for activities
AGENT_APP.on_before_activity(_load_history)
AGENT_APP.on_after_activity(_save_history)


async def _welcome_and_help(context: TurnContext, state: PersistentToolUsingAgentState):
    await context.send_activity(
        "Hello! I am a persistent Agent with web search capabilities. "
        "Ask me questions like 'What is the weather in London?' or 'What is the current time?'"
    )
    state.history.append({"role": "assistant", "content": "Hello! I am a persistent Agent with web search capabilities. Ask me questions like 'What is the weather in London?' or 'What is the current time?'"})


AGENT_APP.conversation_update("membersAdded")(_welcome_and_help)
AGENT_APP.message("/help")(_welcome_and_help)


@AGENT_APP.activity("message")
async def on_message_with_tools_persistent(context: TurnContext, state: PersistentToolUsingAgentState):
    user_message = context.activity.text
    state.history.append({"role": "user", "content": user_message})

    messages_for_llm = [
        {"role": "system", "content": "You are a helpful assistant. Use tools if necessary to answer questions."},
        *state.history
    ]

    try:
        response = await aoai_client.chat.completions.create(
            model=AZURE_OPENAI_DEPLOYMENT_NAME,
            messages=messages_for_llm,
            tools=get_tool_schemas(),
            tool_choice="auto"
        )

        assistant_message = response.choices[0].message
        if assistant_message.content:
            state.history.append({"role": "assistant", "content": assistant_message.content})
        
        if assistant_message.tool_calls:
            for tool_call in assistant_message.tool_calls:
                function_name = tool_call.function.name
                arguments = tool_call.function.arguments

                print(f"Agent wants to call tool: {function_name} with args: {arguments}")

                tool_output = None
                if function_name == "web_search":
                    query = json.loads(arguments).get("query")
                    tool_output = await SearchTool.web_search(query)
                elif function_name == "get_current_time":
                    tool_output = await SearchTool.get_current_time()
                else:
                    tool_output = f"Unknown tool: {function_name}"

                state.history.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": str(tool_output)
                })

                final_response = await aoai_client.chat.completions.create(
                    model=AZURE_OPENAI_DEPLOYMENT_NAME,
                    messages=state.history
                )
                final_assistant_message_content = final_response.choices[0].message.content
                state.history.append({"role": "assistant", "content": final_assistant_message_content})
                await context.send_activity(final_assistant_message_content)

        else:
            await context.send_activity(assistant_message.content)

    except Exception as e:
        print(f"Error in persistent tool-using agent: {e}")
        error_msg = "An error occurred while processing your request with persistent memory."
        await context.send_activity(error_msg)
        state.history.append({"role": "assistant", "content": error_msg})


# --- Server Setup (same as before) ---
def start_server(agent_application: AgentApplication):
    async def entry_point(req: Request) -> Response:
        agent: AgentApplication = req.app["agent_app"]
        adapter: CloudAdapter = req.app["adapter"]
        return await start_agent_process(req, agent, adapter)

    APP = Application()
    APP.router.add_post("/api/messages", entry_point)
    APP.router.add_get("/api/messages", lambda _: Response(status=200))
    APP["agent_app"] = agent_application
    APP["adapter"] = agent_application.adapter
    try:
        port = int(os.environ.get("PORT", 3978))
        print(f"======== Running on http://localhost:{port} ========")
        print("(Press CTRL+C to quit)")
        run_app(APP, host="localhost", port=port)
    except Exception as error:
        print(f"Error starting server: {error}")
        raise error

if __name__ == "__main__":
    start_server(AGENT_APP)

# tools.py (Updated with content-specific tools)
import asyncio
from typing import Dict, Any, List

class SearchTool:
    @staticmethod
    async def web_search(query: str) -> List[str]:
        """
        Performs a simulated web search and returns relevant snippets.
        Used by the writer for research.
        """
        print(f"Executing web search for: {query}")
        await asyncio.sleep(0.5)
        if "AI agent framework" in query.lower():
            return ["Microsoft Agent Framework unifies Semantic Kernel and AutoGen.", "LangChain is popular for LLM orchestration."]
        return [f"Search results for '{query}': Example content related to the query."]

class ContentReviewTool:
    @staticmethod
    async def check_grammar(text: str) -> str:
        """
        Checks the grammar of the provided text and suggests corrections.
        Used by the editor.
        """
        print(f"Executing grammar check for text: {text[:50]}...")
        await asyncio.sleep(0.3)
        if " an example" in text:
            return text.replace(" an example", " a clear example") + "\n(Suggested: 'a clear example' instead of 'an example')"
        return text + "\n(Grammar check: No significant issues found.)"

    @staticmethod
    async def analyze_sentiment(text: str) -> str:
        """
        Analyzes the sentiment of the provided text.
        Used by the editor.
        """
        print(f"Executing sentiment analysis for text: {text[:50]}...")
        await asyncio.sleep(0.3)
        if "negative" in text.lower() or "bad" in text.lower():
            return "Sentiment: Negative"
        return "Sentiment: Neutral/Positive"

# multi_agent_workflow.py
import os
import asyncio
import json
from dotenv import load_dotenv

load_dotenv()

from microsoft.agents.hosting.core import AgentApplication, TurnState, TurnContext, MemoryStorage
from microsoft.agents.hosting.aiohttp import CloudAdapter, start_agent_process
from aiohttp.web import Request, Response, Application, run_app
from openai import AzureOpenAI

from tools import SearchTool, ContentReviewTool # Our updated tools

# --- Configuration (same as before) ---
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME")

if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_DEPLOYMENT_NAME]):
    raise ValueError("Missing Azure OpenAI environment variables. Please check your .env file.")

try:
    aoai_client = AzureOpenAI(
        api_key=AZURE_OPENAI_API_KEY,
        api_version="2024-02-15-preview",
        azure_endpoint=AZURE_OPENAI_ENDPOINT,
    )
    print("Successfully connected to Azure OpenAI.")
except Exception as e:
    print(f"Error connecting to Azure OpenAI: {e}")
    exit(1)

class AzureOpenAIChatClient:
    def __init__(self, client, deployment_name):
        self.client = client
        self.deployment_name = deployment_name

    async def complete(self, messages: list[dict], **kwargs) -> str:
        try:
            response = await self.client.chat.completions.create(
                model=self.deployment_name,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error calling Azure OpenAI: {e}")
            return "An error occurred while processing your request."

my_chat_client = AzureOpenAIChatClient(aoai_client, AZURE_OPENAI_DEPLOYMENT_NAME)


# --- Tool Definitions for LLM ---
def get_writer_tool_schemas():
    return [
        {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "Performs a web search to gather information for writing.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query."
                        }
                    },
                    "required": ["query"]
                }
            }
        }
    ]

def get_editor_tool_schemas():
    return [
        {
            "type": "function",
            "function": {
                "name": "check_grammar",
                "description": "Checks the grammar of the provided text and suggests corrections.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "The text to check."
                        }
                    },
                    "required": ["text"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "analyze_sentiment",
                "description": "Analyzes the sentiment of the provided text.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "text": {
                            "type": "string",
                            "description": "The text to analyze."
                        }
                    },
                    "required": ["text"]
                }
            }
        }
    ]

# --- Base Agent Class for MAF-like structure ---
class BaseMAFAgent:
    def __init__(self, name: str, instructions: str, tool_schemas: List[Dict[str, Any]] = None):
        self.name = name
        self.instructions = instructions
        self.tool_schemas = tool_schemas if tool_schemas is not None else []
        self.aoai_client = aoai_client # Using global client for simplicity

    async def run_step(self, messages: List[Dict[str, Any]]) -> Dict[str, Any]:
        """Runs a single step of the agent, potentially calling tools."""
        current_messages = [{"role": "system", "content": self.instructions}] + messages
        try:
            response = await self.aoai_client.chat.completions.create(
                model=AZURE_OPENAI_DEPLOYMENT_NAME,
                messages=current_messages,
                tools=self.tool_schemas,
                tool_choice="auto"
            )
            assistant_message = response.choices[0].message
            
            output = {"content": assistant_message.content, "tool_calls": []}
            if assistant_message.tool_calls:
                output["tool_calls"] = assistant_message.tool_calls
            return output
        except Exception as e:
            print(f"Error in {self.name} agent: {e}")
            return {"content": f"An error occurred in {self.name}.", "tool_calls": []}

    async def execute_tool_call(self, tool_call) -> str:
        function_name = tool_call.function.name
        arguments = json.loads(tool_call.function.arguments)

        print(f"  {self.name} executing tool: {function_name} with args: {arguments}")

        tool_output = f"Unknown tool: {function_name}"
        if function_name == "web_search":
            query = arguments.get("query")
            tool_output = await SearchTool.web_search(query)
        elif function_name == "check_grammar":
            text = arguments.get("text")
            tool_output = await ContentReviewTool.check_grammar(text)
        elif function_name == "analyze_sentiment":
            text = arguments.get("text")
            tool_output = await ContentReviewTool.analyze_sentiment(text)
        
        return str(tool_output)

# --- Define specific Agents ---
writer_agent = BaseMAFAgent(
    name="Writer",
    instructions="You are a creative writer. Your goal is to write engaging content based on user requests. Use the web_search tool for research.",
    tool_schemas=get_writer_tool_schemas()
)

editor_agent = BaseMAFAgent(
    name="Editor",
    instructions="You are an expert editor. Your goal is to review and improve written content, focusing on grammar, engagement, and tone. Use check_grammar and analyze_sentiment tools.",
    tool_schemas=get_editor_tool_schemas()
)

# --- Workflow Logic (Simulated by manual chaining for Python MAF SDK) ---
# In MAF Python SDK, this would typically involve an AgentWorkflowBuilder or similar
# for declarative workflow definition. For now, we simulate the sequential flow.

class ContentWorkflowState(TurnState):
    history: List[Dict[str, Any]] = []
    current_content: str = ""
    user_request: str = ""

# A single AgentApplication to manage the workflow
WORKFLOW_AGENT_APP = AgentApplication[ContentWorkflowState](
    storage=MemoryStorage(), # Using in-memory for this example
    adapter=CloudAdapter()
)

async def _workflow_welcome(context: TurnContext, state: ContentWorkflowState):
    await context.send_activity(
        "Welcome to the Content Creation Workflow! I have a Writer and an Editor agent. "
        "Tell me what content you want to create (e.g., 'Write an article about AI agent frameworks')."
    )
    state.history.append({"role": "assistant", "content": "Welcome to the Content Creation Workflow!"})

WORKFLOW_AGENT_APP.conversation_update("membersAdded")(_workflow_welcome)
WORKFLOW_AGENT_APP.message("/help")(_workflow_welcome)

@WORKFLOW_AGENT_APP.activity("message")
async def content_creation_workflow_handler(context: TurnContext, state: ContentWorkflowState):
    user_message = context.activity.text
    state.user_request = user_message
    state.history.append({"role": "user", "content": user_message})
    await context.send_activity(f"Starting content creation for: '{user_message}'...")

    # --- Step 1: Writer Agent writes the content ---
    await context.send_activity("Writer Agent is working...")
    writer_messages = [{"role": "user", "content": f"Write content based on the request: {user_message}"}]
    
    writer_output = await writer_agent.run_step(writer_messages)
    content_draft = writer_output["content"]
    
    # Handle writer's tool calls
    if writer_output["tool_calls"]:
        for tool_call in writer_output["tool_calls"]:
            tool_result = await writer_agent.execute_tool_call(tool_call)
            # Send tool result back to writer for continued reasoning
            writer_messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": tool_result})
            re_writer_output = await writer_agent.run_step(writer_messages)
            content_draft = re_writer_output["content"] # Use the refined content
    
    await context.send_activity(f"Writer Draft:\n```\n{content_draft}\n```")
    state.current_content = content_draft
    state.history.append({"role": "assistant", "content": f"Writer produced draft: {content_draft}"})

    # --- Step 2: Editor Agent reviews and refines ---
    await context.send_activity("Editor Agent is reviewing...")
    editor_messages = [{"role": "user", "content": f"Review and improve this content: {state.current_content}"}]
    
    editor_output = await editor_agent.run_step(editor_messages)
    final_content = editor_output["content"]

    # Handle editor's tool calls
    if editor_output["tool_calls"]:
        for tool_call in editor_output["tool_calls"]:
            tool_result = await editor_agent.execute_tool_call(tool_call)
            editor_messages.append({"role": "tool", "tool_call_id": tool_call.id, "content": tool_result})
            re_editor_output = await editor_agent.run_step(editor_messages)
            final_content = re_editor_output["content"] # Use the refined content
    
    await context.send_activity(f"Editor Final Version:\n```\n{final_content}\n```")
    state.history.append({"role": "assistant", "content": f"Editor finalized content: {final_content}"})
    await context.send_activity("Content creation workflow completed!")


# --- Server Setup (same as before) ---
def start_server(agent_application: AgentApplication):
    async def entry_point(req: Request) -> Response:
        agent: AgentApplication = req.app["agent_app"]
        adapter: CloudAdapter = req.app["adapter"]
        return await start_agent_process(req, agent, adapter)

    APP = Application()
    APP.router.add_post("/api/messages", entry_point)
    APP.router.add_get("/api/messages", lambda _: Response(status=200))
    APP["agent_app"] = agent_application
    APP["adapter"] = agent_application.adapter
    try:
        port = int(os.environ.get("PORT", 3978))
        print(f"======== Running on http://localhost:{port} ========")
        print("(Press CTRL+C to quit)")
        run_app(APP, host="localhost", port=port)
    except Exception as error:
        print(f"Error starting server: {error}")
        raise error

if __name__ == "__main__":
    start_server(WORKFLOW_AGENT_APP)

# secured_agent.py (Illustrates JWT Middleware usage, building on redis_agent.py)
import os
import asyncio
import json
from dotenv import load_dotenv
import redis.asyncio as redis

load_dotenv()

from microsoft.agents.hosting.core import AgentApplication, TurnState, TurnContext, MemoryStorage
from microsoft.agents.hosting.aiohttp import CloudAdapter, start_agent_process, jwt_authorization_middleware
from aiohttp.web import Request, Response, Application, run_app, HTTPUnauthorized, HTTPException
from openai import AzureOpenAI
from azure.identity import DefaultAzureCredential, ClientSecretCredential

from tools import SearchTool

# --- Configuration (same as before) ---
AZURE_OPENAI_API_KEY = os.getenv("AZURE_OPENAI_API_KEY")
AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_DEPLOYMENT_NAME = os.getenv("AZURE_OPENAI_DEPLOYMENT_NAME")

REDIS_HOST = os.getenv("REDIS_HOST", "localhost")
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD")

# --- Microsoft Entra ID Configuration (for validation) ---
# For actual JWT validation, you'd configure the middleware with issuer, audience, etc.
# This is typically done via environment variables or a configuration object.
# These placeholders are for illustrating the *presence* of auth.
# REAL-WORLD: You would register an application in Azure Portal (App Registrations)
# and get these values.
# AZURE_TENANT_ID = os.getenv("AZURE_TENANT_ID")
# AZURE_CLIENT_ID = os.getenv("AZURE_CLIENT_ID")
# AZURE_CLIENT_SECRET = os.getenv("AZURE_CLIENT_SECRET") # If using client secret flow

if not all([AZURE_OPENAI_API_KEY, AZURE_OPENAI_ENDPOINT, AZURE_OPENAI_DEPLOYMENT_NAME]):
    raise ValueError("Missing Azure OpenAI environment variables. Please check your .env file.")

# AOAI Client and Redis Client initialization (same as redis_agent.py)
try:
    aoai_client = AzureOpenAI(
        api_key=AZURE_OPENAI_API_KEY,
        api_version="2024-02-15-preview",
        azure_endpoint=AZURE_OPENAI_ENDPOINT,
    )
    print("Successfully connected to Azure OpenAI.")
except Exception as e:
    print(f"Error connecting to Azure OpenAI: {e}")
    exit(1)

try:
    redis_client = redis.Redis(
        host=REDIS_HOST,
        port=REDIS_PORT,
        password=REDIS_PASSWORD,
        decode_responses=True
    )
    async def _test_redis_connection():
        await redis_client.ping()
        print("Successfully connected to Redis.")
    asyncio.run(_test_redis_connection())
except Exception as e:
    print(f"Error connecting to Redis: {e}")
    redis_client = None


class AzureOpenAIChatClient:
    def __init__(self, client, deployment_name):
        self.client = client
        self.deployment_name = deployment_name

    async def complete(self, messages: list[dict], **kwargs) -> str:
        try:
            response = await self.client.chat.completions.create(
                model=self.deployment_name,
                messages=messages,
                **kwargs
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Error calling Azure OpenAI: {e}")
            return "An error occurred while processing your request."

my_chat_client = AzureOpenAIChatClient(aoai_client, AZURE_OPENAI_DEPLOYMENT_NAME)

def get_tool_schemas():
    return [
        {
            "type": "function",
            "function": {
                "name": "web_search",
                "description": "Performs a web search to find relevant information.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "query": {
                            "type": "string",
                            "description": "The search query."
                        }
                    },
                    "required": ["query"]
                }
            }
        }
    ]

class RedisChatHistoryStorage(MemoryStorage):
    def __init__(self, redis_client: redis.Redis, prefix: str = "chat_history:"):
        self.redis_client = redis_client
        self.prefix = prefix

    async def _get_key(self, conversation_id: str) -> str:
        return f"{self.prefix}{conversation_id}"

    async def read(self, conversation_id: str) -> dict:
        key = await self._get_key(conversation_id)
        raw_history = await self.redis_client.get(key)
        if raw_history:
            return json.loads(raw_history)
        return {"history": []}

    async def write(self, conversation_id: str, data: dict):
        key = await self._get_key(conversation_id)
        await self.redis_client.set(key, json.dumps(data))

class PersistentToolUsingAgentState(TurnState):
    history: List[Dict[str, Any]] = []
    conversation_id: str = "default_conversation"

storage_instance = RedisChatHistoryStorage(redis_client) if redis_client else MemoryStorage()

# --- Agent Application (same as redis_agent.py) ---
SECURE_AGENT_APP = AgentApplication[PersistentToolUsingAgentState](
    storage=storage_instance,
    adapter=CloudAdapter()
)

async def _load_history(context: TurnContext, state: PersistentToolUsingAgentState):
    if context.activity and context.activity.conversation and context.activity.conversation.id:
        state.conversation_id = context.activity.conversation.id
    if redis_client:
        stored_data = await storage_instance.read(state.conversation_id)
        state.history = stored_data.get("history", [])
        print(f"Loaded history for {state.conversation_id}: {len(state.history)} items")

async def _save_history(context: TurnContext, state: PersistentToolUsingAgentState):
    if redis_client and state.conversation_id:
        await storage_instance.write(state.conversation_id, {"history": state.history})
        print(f"Saved history for {state.conversation_id}: {len(state.history)} items")

SECURE_AGENT_APP.on_before_activity(_load_history)
SECURE_AGENT_APP.on_after_activity(_save_history)

async def _secure_welcome_and_help(context: TurnContext, state: PersistentToolUsingAgentState):
    await context.send_activity(
        "Hello! I am a **secured** persistent Agent with web search capabilities. "
        "Your access is authenticated via JWT."
    )
    state.history.append({"role": "assistant", "content": "Hello! I am a secured persistent Agent."})

SECURE_AGENT_APP.conversation_update("membersAdded")(_secure_welcome_and_help)
SECURE_AGENT_APP.message("/help")(_secure_welcome_and_help)

@SECURE_AGENT_APP.activity("message")
async def on_secure_message(context: TurnContext, state: PersistentToolUsingAgentState):
    # This handler is the same as in redis_agent.py, but now it's protected
    user_message = context.activity.text
    state.history.append({"role": "user", "content": user_message})

    messages_for_llm = [
        {"role": "system", "content": "You are a helpful assistant. Use tools if necessary to answer questions."},
        *state.history
    ]

    try:
        response = await aoai_client.chat.completions.create(
            model=AZURE_OPENAI_DEPLOYMENT_NAME,
            messages=messages_for_llm,
            tools=get_tool_schemas(),
            tool_choice="auto"
        )

        assistant_message = response.choices[0].message
        if assistant_message.content:
            state.history.append({"role": "assistant", "content": assistant_message.content})
        
        if assistant_message.tool_calls:
            for tool_call in assistant_message.tool_calls:
                function_name = tool_call.function.name
                arguments = json.loads(tool_call.function.arguments)

                print(f"Agent wants to call tool: {function_name} with args: {arguments}")

                tool_output = None
                if function_name == "web_search":
                    query = arguments.get("query")
                    tool_output = await SearchTool.web_search(query)
                else:
                    tool_output = f"Unknown tool: {function_name}"

                state.history.append({
                    "tool_call_id": tool_call.id,
                    "role": "tool",
                    "name": function_name,
                    "content": str(tool_output)
                })

                final_response = await aoai_client.chat.completions.create(
                    model=AZURE_OPENAI_DEPLOYMENT_NAME,
                    messages=state.history
                )
                final_assistant_message_content = final_response.choices[0].message.content
                state.history.append({"role": "assistant", "content": final_assistant_message_content})
                await context.send_activity(final_assistant_message_content)

        else:
            await context.send_activity(assistant_message.content)

    except Exception as e:
        print(f"Error in secured agent: {e}")
        error_msg = "An error occurred in the secured agent."
        await context.send_activity(error_msg)
        state.history.append({"role": "assistant", "content": error_msg})


# --- Server Setup with JWT Middleware ---
# The jwt_authorization_middleware needs configuration.
# For simplicity, we'll use a dummy config. In a real scenario, this would involve
# proper Azure AD app registration details for token validation.
from microsoft.agents.hosting.core import AgentAuthConfiguration

class DummyAgentAuthConfiguration(AgentAuthConfiguration):
    def __init__(self):
        # In a real app, these would come from env vars or app settings
        self.tenant_id = "your-azure-ad-tenant-id"
        self.client_id = "your-azure-ad-client-id"
        self.authority = f"https://login.microsoftonline.com/{self.tenant_id}"
        self.audience = [f"api://{self.client_id}"] # The audience of your agent's API
        self.is_anonymous_auth_enabled = False # Require authentication

DUMMY_AUTH_CONFIG = DummyAgentAuthConfiguration()

def start_server_with_auth(agent_application: AgentApplication, auth_config: AgentAuthConfiguration):
    async def entry_point(req: Request) -> Response:
        # After middleware, req['user_id'] or similar would be populated
        # This is where you might add custom authorization checks if needed
        agent: AgentApplication = req.app["agent_app"]
        adapter: CloudAdapter = req.app["adapter"]
        return await start_agent_process(req, agent, adapter)

    # Apply the JWT middleware
    # In a real app, this middleware would check for a valid JWT token in the Authorization header.
    APP = Application(middlewares=[jwt_authorization_middleware(auth_config=auth_config)])
    APP.router.add_post("/api/messages", entry_point)
    APP.router.add_get("/api/messages", lambda _: Response(status=200)) # Health check
    APP["agent_app"] = agent_application
    APP["adapter"] = agent_application.adapter
    APP["agent_configuration"] = auth_config # Store auth config for middleware
    try:
        port = int(os.environ.get("PORT", 3978))
        print(f"======== Running on http://localhost:{port} (Secured with JWT) ========")
        print("(Press CTRL+C to quit)")
        run_app(APP, host="localhost", port=port)
    except Exception as error:
        print(f"Error starting secured server: {error}")
        raise error

if __name__ == "__main__":
    # Note: To test this locally with actual JWT, you'd need to acquire a token
    # (e.g., using MSAL Python library or Azure CLI) and include it in your client's
    # Authorization header. teamsapptester usually handles anonymous or specific auth.
    # For a quick test to see the auth barrier, remove `is_anonymous_auth_enabled` from
    # DummyAgentAuthConfiguration and send a request without a token - you should get 401.
    start_server_with_auth(SECURE_AGENT_APP, DUMMY_AUTH_CONFIG)