from agents import Agent, OpenAIChatCompletionsModel, AsyncOpenAI, ModelSettings
from agents.mcp import MCPServerStreamableHttp
from backend.agent.sys_prompt import get_system_prompt
from backend.agent.hooks import LegalAssistantHooks
from configs import (
    LITELLM_BASE_URL, 
    LITELLM_API_KEY, 
    LLM_TIMEOUT,
    MCP_SERVER_URL,
    AGENT_TEMPERATURE,
    DEFAULT_MODEL
)

def make_client() -> AsyncOpenAI:
    """Async client pointing to LiteLLM proxy."""
    return AsyncOpenAI (
        base_url=LITELLM_BASE_URL,
        api_key=LITELLM_API_KEY,
        timeout=LLM_TIMEOUT,
        max_retries=0
    )

def make_mcp_server() -> MCPServerStreamableHttp:
    """MCP server exposing Slovak Justice API tools."""
    return MCPServerStreamableHttp(
        name="Slovak Justice API",
        params={"url": MCP_SERVER_URL},
        cache_tools_list=True
    )

def build_agent(mcp_server: MCPServerStreamableHttp, model_name: str = DEFAULT_MODEL) -> Agent:
    """
    Build and return a Legal AI Assistant agent.

    Args:
        model_name: Model identifier routed through LiteLLM.

    Returns:
        Configured Agent ready to run.
    """
    
    return Agent (
        name="Legal AI Assistant",
        instructions=get_system_prompt(),
        model=OpenAIChatCompletionsModel(
            model=model_name, 
            openai_client=make_client()
        ),
        model_settings=ModelSettings(
            temperature=AGENT_TEMPERATURE,
            tool_choice="auto",
            parallel_tool_calls=False,
            include_usage=True
        ),
        mcp_servers=[mcp_server],
        tool_use_behavior="run_llm_again",
        reset_tool_choice=True,
        hooks=LegalAssistantHooks()
    )