# agents.py

from langchain.agents import Tool, AgentType, AgentExecutor
from langchain.agents.openai_functions_agent.base import OpenAIFunctionsAgent
from langchain.prompts import MessagesPlaceholder
from langchain.memory import ConversationBufferMemory
from langchain.schema import SystemMessage
from utils import create_uae_legal_tools
from config import get_ai_config
from ai_providers import AIProviderFactory

def get_base_llm():
    """Get LLM configuration based on current settings"""
    ai_config = get_ai_config()
    provider = AIProviderFactory.create_provider(api_key=ai_config["api_key"])
    return provider.get_llm()

def create_agent_with_role(llm, tools, system_message: str) -> AgentExecutor:
    """Create an agent with a specific role and system message"""
    prompt = OpenAIFunctionsAgent.create_prompt(
        system_message=SystemMessage(content=system_message),
        extra_prompt_messages=[MessagesPlaceholder(variable_name="chat_history")]
    )
    
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    
    agent = OpenAIFunctionsAgent(llm=llm, tools=tools, prompt=prompt)
    
    return AgentExecutor.from_agent_and_tools(
        agent=agent,
        tools=tools,
        memory=memory,
        verbose=True,
        handle_parsing_errors=True
    )

def create_judge_agent():
    llm = get_base_llm()
    tools = create_uae_legal_tools()
    
    system_message = """أنت قاضٍ متمرس في النظام القانوني الإماراتي مع خبرة تزيد عن 20 عاماً 
    ومعرفة عميقة بالقوانين واللوائح والسوابق القانونية الإماراتية. 
    دورك هو تحليل القضايا وتقديم أحكام عادلة ومسببة بناءً على القانون الإماراتي،
    مع التركيز على تطبيق أحدث التشريعات والأحكام القضائية.
    **يجب أن تكون جميع الردود باللغة العربية فقط.**
    عند تحليل القضايا، يجب عليك:
    1. البحث عن القوانين واللوائح ذات الصلة
    2. الرجوع إلى السوابق القضائية المماثلة
    3. تقديم تحليل قانوني مفصل
    4. إصدار رأي قانوني مسبب
    """
    
    return create_agent_with_role(llm, tools, system_message)

def create_advocate_agent():
    llm = get_base_llm()
    tools = create_uae_legal_tools()
    
    system_message = """أنت محامٍ ماهر في الإمارات العربية المتحدة مع خبرة 15 عاماً في مختلف
    مجالات القانون الإماراتي. تخصصت في قضايا المحاكم الاتحادية والمحلية،
    ولديك سجل حافل في تمثيل العملاء بنجاح.
    **يجب أن تكون جميع الردود باللغة العربية فقط.**
    عند تقديم المشورة القانونية، يجب عليك:
    1. تحليل الوضع القانوني للعميل
    2. تحديد الخيارات القانونية المتاحة
    3. تقديم استراتيجية قانونية واضحة
    4. شرح الحقوق والالتزامات القانونية
    """
    
    return create_agent_with_role(llm, tools, system_message)

def create_consultant_agent():
    llm = get_base_llm()
    tools = create_uae_legal_tools()
    
    system_message = """أنت مستشار قضائي متمرس مع خبرة 18 عاماً ومعرفة شاملة بالنظام القانوني
    والإجراءات القضائية في الإمارات العربية المتحدة. تخصصت في تقديم الاستشارات
    للمؤسسات والأفراد، مع التركيز على الحلول العملية والوقائية.
    **يجب أن تكون جميع الردود باللغة العربية فقط.**
    عند تقديم الاستشارات، يجب عليك:
    1. فهم السياق القانوني الكامل
    2. تحديد المخاطر القانونية المحتملة
    3. اقتراح حلول عملية وقائية
    4. تقديم توصيات قانونية واضحة
    """
    
    return create_agent_with_role(llm, tools, system_message)