import sys
import os

# Add the current directory to path for imports
sys.path.append(os.getcwd())

from core.agent import NyayaAgent
from tools.search_tool import search_indian_kanoon, get_act_details, search_legal_documents
from tools.legal_tools import draft_rti_query, draft_legal_notice, analyze_case_summary, suggest_sc_followup
from tools.compliance import audit_privacy_policy
from tools.law_bridge import ipc_to_bns_lookup
from google.genai import types

def main():
    print("--- Nyaya Agent: Indian Legal Assistant ---")
    print("Type 'exit' to quit.\n")
    
    agent = NyayaAgent()
    
    # Define tools for function calling
    tools = [
        search_indian_kanoon, 
        get_act_details, 
        search_legal_documents,
        draft_rti_query,
        draft_legal_notice,
        analyze_case_summary,
        suggest_sc_followup,
        audit_privacy_policy,
        ipc_to_bns_lookup
    ]
    
    config = types.GenerateContentConfig(
        system_instruction=agent.system_instruction,
        tools=tools,
        temperature=0.2,
    )

    while True:
        user_input = input("You: ")
        if user_input.lower() in ['exit', 'quit']:
            break
            
        try:
            # Using the agent's wrapper which includes retry logic
            response_text = agent.generate_response(user_input, config=config)
            print(f"\nNyaya Agent: {response_text}")
            
        except Exception as e:
            print(f"An error occurred: {e}")

if __name__ == "__main__":
    main()