from app.services.memory import memory_service
from app.services.guardrails import guardrail_service, GuardrailViolation
from app.core.swarm_client import swarm_client
from app.agents.gatekeeper import gatekeeper_agent
from typing import Dict

class OrchestratorService:
    async def process_message(self, user_id: str, message: str) -> dict:
        # 1. Safety Check (Input)
        is_safe = await guardrail_service.check_input_safety(message)
        if not is_safe:
            return {
                "role": "assistant",
                "content": "I sense you are in deep distress. Please contact a professional or emergency services immediately. You are valuable and loved."
            }

        # 2. Retrieve Episodic Context (Phase 2)
        episodes = await memory_service.retrieve_episodes(user_id, message, limit=2)
        episodic_context = ""
        if episodes:
            insights = [f"- {e['timestamp']}: {e['insight']}" for e in episodes]
            episodic_context = "\nRelevant past insights:\n" + "\n".join(insights)

        # 3. Run Swarm
        messages = [{"role": "user", "content": message}]
        
        # Inject episodic memory into context variables for agents to see
        context = {
            "user_id": user_id,
            "episodic_memory": episodic_context
        }
        
        swarm_response = await swarm_client.run(
            agent=gatekeeper_agent,
            messages=messages,
            context_variables=context
        )
        
        # The final message in the swarm history is the result
        response_text = swarm_response.messages[-1]["content"]

        # 3. Final Safety Validation of Output
        try:
             response_text = await guardrail_service.validate_response(response_text)
        except GuardrailViolation as e:
             print(f"Safety Violation blocked: {e}")
             response_text = "I apologize, but I cannot provide that response as it violates my safety guidelines regarding authority claims."
        except GuardrailViolation:
            response_text = "I apologize, but I cannot complete that response as it violates my safety guidelines."

        return {
            "role": "assistant",
            "content": response_text,
            "agent": swarm_response.agent.name if swarm_response.agent else "Unknown",
            "trace": swarm_response.trace
        }

orchestrator = OrchestratorService()