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_mena_law(query: str, jurisdiction: str = "EG") -> str:
    """Search Egyptian (EG) or Saudi (SA) legislation and court precedents."""
    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}, "top_k": 5},
        )
        return str(res.json().get("results", []))

@tool
def reason_mena_law(question: str, jurisdiction: str = "EG") -> str:
    """Ask a complex legal question to receive an authoritative answer with verified citations."""
    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}},
            timeout=30.0,
        )
        data = res.json()
        return f"Answer: {data.get('answer')}\nCitations: {data.get('citations')}"

# Initialize Agent
llm = ChatOpenAI(model="gpt-4o")
tools = [search_mena_law, reason_mena_law]
agent = create_react_agent(llm, tools)

# Run Query
response = agent.invoke({
    "messages": [("user", "ما هي شروط فسخ عقد المقاولة في القانون المصري عند تأخر المقاول؟")]
})
print(response["messages"][-1].content)

import { tool } from 'ai';
import { z } from 'zod';

export const dikeSearchTool = tool({
  description: 'Search MENA legal statutes, court rulings, and regulations across Egypt and Saudi Arabia.',
  parameters: z.object({
    query: z.string().describe('The Arabic or English legal query'),
    jurisdiction: z.enum(['EG', 'SA', 'AE']).default('EG'),
  }),
  execute: async ({ query, jurisdiction }) => {
    const res = await fetch('https://api.dike.it.com/v1/search', {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${process.env.DIKE_API_KEY}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        query,
        filters: { jurisdiction },
        top_k: 5,
        rerank: true,
      }),
    });

    const data = await res.json();
    return data.results;
  },
});

export const dikeReasonTool = tool({
  description: 'Run grounded legal reasoning with verified citation footnotes.',
  parameters: z.object({
    question: z.string().describe('The legal inquiry in Arabic or English'),
    jurisdiction: z.enum(['EG', 'SA', 'AE']).default('EG'),
  }),
  execute: async ({ question, jurisdiction }) => {
    const res = await fetch('https://api.dike.it.com/v1/reason', {
      method: 'POST',
      headers: {
        'Authorization': `Bearer ${process.env.DIKE_API_KEY}`,
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({
        question,
        filters: { jurisdiction },
        language: 'ar',
      }),
    });

    return await res.json();
  },
});

import os
import httpx
from llama_index.core.tools import FunctionTool
from llama_index.core.agent import ReActAgent
from llama_index.llms.openai import OpenAI

DIKE_API_KEY = os.environ["DIKE_API_KEY"]

def legal_graph_lookup(canonical_id: str) -> dict:
    """Traverse the citation graph for a law to find amendments and repeals."""
    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": "both", "depth": 1},
        )
        return res.json()

graph_tool = FunctionTool.from_defaults(fn=legal_graph_lookup)
agent = ReActAgent.from_tools([graph_tool], llm=OpenAI(model="gpt-4o"))

response = agent.chat("ما هي القوانين التي قامت بتعديل القانون المدني المصري رقم 131 لسنة 1948؟")
print(response)