from langchain_core.pydantic_v1 import BaseModel, Field, root_validator
from langchain.tools import tool
import dns.resolver, dns.reversename
import validators
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
from hw1.app import get_rag_chain

"""
Tools to be run by my custom agent.

So far, these are the same tooling provided by the example in https://github.com/wu4f/cs410g-src/blob/main/04_Agents/07_tools_custom_agent.py,
But I will be updating this to build my own tooling
"""

class LookupNameInput(BaseModel):
    hostname: str = Field(description="Should be a hostname such as www.google.com")
    @root_validator
    def is_dns_address(cls, values: dict[str,any]) -> str:
        if validators.domain(values.get("hostname")):
            return values
        raise ValueError("Malformed hostname")
class LookupIPInput(BaseModel):
    address: str = Field(description="Should be an IP address such as 208.91.197.27 or 143.95.239.83")
    @root_validator
    def is_ip_address(cls, values: dict[str,any]) -> str:
        if validators.ip_address.ipv4(values.get("address")):
            return values
        raise ValueError("Malformed IP address")
    
class KSPTool(BaseModel):
    query: str = Field(description="should be a kerbal space program (ksp) related query")

@tool("kerbal_space_program_ksp_information", args_schema=KSPTool, return_direct=False)
def search_ksp(query:str) -> str:
    """Given a query about kerbal space program (ksp), it will send the query to the KSP rag applciation"""
    return get_rag_chain().invoke(query)

@tool("lookup_name",args_schema=LookupNameInput, return_direct=False)
def lookup_name(hostname):
    """Given a DNS hostname, it will return its IPv4 addresses"""
    result = dns.resolver.resolve(hostname, 'A')
    res = [ r.to_text() for r in result ]
    return res[0]


@tool("lookup_ip", args_schema=LookupIPInput, return_direct=False)
def lookup_ip(address):
    """Given an IP address, returns names associated with it"""
    n = dns.reversename.from_address(address)
    result = dns.resolver.resolve(n, 'PTR')
    res = [ r.to_text() for r in result ]
    return res[0]