import os
import httpx
from langchain_core.tools import tool
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import create_react_agent

DIKE_API_KEY = os.environ["DIKE_API_KEY"]

@tool
def search_statute(query: str, jurisdiction: str = "SA") -> dict:
    """Search Saudi (SA) or Egyptian (EG) legislation and regulations."""
    with httpx.Client(base_url="https://api.dike.it.com") as client:
        res = client.post(
            "/v1/search",
            headers={"Authorization": f"Bearer {DIKE_API_KEY}"},
            json={
                "query": query,
                "filters": {
                    "jurisdiction": jurisdiction,
                    "status": ["in_force", "amended"],
                },
                "top_k": 3,
                "rerank": True,
            },
            timeout=15.0,
        )
        return res.json()

@tool
def check_amendment_history(canonical_id: str) -> dict:
    """Audit whether a specific statute or article was subsequently amended or repealed."""
    with httpx.Client(base_url="https://api.dike.it.com") as client:
        res = client.post(
            "/v1/graph/traverse",
            headers={"Authorization": f"Bearer {DIKE_API_KEY}"},
            json={
                "canonical_id": canonical_id,
                "direction": "in",
                "edge_types": ["amends", "repeals"],
            },
            timeout=15.0,
        )
        return res.json()

@tool
def generate_grounded_memo(question: str, jurisdiction: str = "SA") -> dict:
    """Synthesize final audit memo with deterministic hallucination guard and verified footnotes."""
    with httpx.Client(base_url="https://api.dike.it.com") as client:
        res = client.post(
            "/v1/reason",
            headers={"Authorization": f"Bearer {DIKE_API_KEY}"},
            json={
                "question": question,
                "filters": {"jurisdiction": jurisdiction},
                "language": "ar",
            },
            timeout=30.0,
        )
        return res.json()

# Initialize ReAct Agent
tools = [search_statute, check_amendment_history, generate_grounded_memo]
model = ChatOpenAI(model="gpt-4o", temperature=0.0)
mna_agent = create_react_agent(model, tools)

# Example Execution
response = mna_agent.invoke({
    "messages": [
        ("user", "ما هي الشروط النظامية لصحة بند عدم المنافسة في عقد العمل وفقاً للنظام السعودي وتعديلاته الحديثة؟")
    ]
})
print(response["messages"][-1].content)