
import subprocess
import requests
import socket
import time
import os
from livekit.agents import function_tool, RunContext
import psutil

# ==============================
# BASIC NETWORK
# ==============================
@function_tool()
async def ip_information(context: RunContext) -> str:
    return requests.get("https://api.ipify.org").text

@function_tool()
async def list_wifi_networks(context: RunContext) -> str:
    res = subprocess.run(["netsh", "wlan", "show", "networks"], capture_output=True, text=True)
    return res.stdout[:1000]

@function_tool()
async def connect_wifi(context: RunContext, ssid: str) -> str:
    subprocess.run(["netsh", "wlan", "connect", f"name={ssid}"])
    return f"Connecting to {ssid}"

@function_tool()
async def disconnect_wifi(context: RunContext) -> str:
    subprocess.run(["netsh", "wlan", "disconnect"])
    return "Disconnected"

@function_tool()
async def get_saved_wifi_passwords(context: RunContext) -> str:
    try:
        output = subprocess.check_output(["netsh", "wlan", "show", "profiles"]).decode('utf-8', errors="ignore")
        profiles = [line.split(":")[1].strip() for line in output.split('\n') if "All User Profile" in line]
        results = []
        for p in profiles:
            try:
                info = subprocess.check_output(["netsh", "wlan", "show", "profile", p, "key=clear"]).decode('utf-8', errors="ignore")
                key = [l for l in info.split('\n') if "Key Content" in l]
                results.append(f"{p}: {key[0].split(':')[1].strip() if key else 'OPEN'}")
            except: pass
        return "\n".join(results)
    except Exception as e: return f"Error: {e}"

# ==============================
# ADVANCED NETWORK
# ==============================
@function_tool()
async def scan_network_devices(context: RunContext) -> str:
    res = subprocess.run(["arp", "-a"], capture_output=True, text=True)
    return res.stdout[:1000]

@function_tool()
async def get_detailed_network_info(context: RunContext) -> str:
    res = subprocess.run(["ipconfig", "/all"], capture_output=True, text=True)
    return res.stdout[:1500]

@function_tool()
async def get_active_connections(context: RunContext) -> str:
    res = subprocess.run(["netstat", "-ano"], capture_output=True, text=True)
    return res.stdout[:1500]

@function_tool()
async def port_scan(context: RunContext, ip_address: str, ports: str = "80,443,22,21,3389") -> str:
    port_list = [int(p.strip()) for p in ports.split(',')]
    results = []
    for port in port_list:
        sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        sock.settimeout(0.5)
        res = sock.connect_ex((ip_address, port))
        results.append(f"Port {port}: {'OPEN' if res == 0 else 'CLOSED'}")
        sock.close()
    return "\n".join(results)

@function_tool()
async def ping_device(context: RunContext, target: str) -> str:
    res = subprocess.run(["ping", "-n", "4", target], capture_output=True, text=True)
    return res.stdout[:1000]

@function_tool()
async def trace_route(context: RunContext, target: str) -> str:
    res = subprocess.run(["tracert", "-h", "10", target], capture_output=True, text=True)
    return res.stdout[:1000]

@function_tool()
async def wake_on_lan(context: RunContext, mac_address: str) -> str:
    # Basic implementation
    try:
        data = b'\xff' * 6 + bytes.fromhex(mac_address.replace(':', '').replace('-', '')) * 16
        with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
            s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
            s.sendto(data, ('255.255.255.255', 9))
        return "Magic packet sent"
    except Exception as e: return f"Error: {e}"

@function_tool()
async def get_router_info(context: RunContext) -> str:
    return subprocess.run(["ipconfig"], capture_output=True, text=True).stdout[:500]

@function_tool()
async def network_bandwidth_usage(context: RunContext) -> str:
    io = psutil.net_io_counters()
    return f"Sent: {io.bytes_sent/1024/1024:.2f}MB | Recv: {io.bytes_recv/1024/1024:.2f}MB"

@function_tool()
async def flush_dns(context: RunContext) -> str:
    subprocess.run("ipconfig /flushdns", shell=True)
    return "DNS Flushed"

@function_tool()
async def renew_ip_address(context: RunContext) -> str:
    subprocess.run("ipconfig /renew", shell=True)
    return "IP Renewed"

@function_tool()
async def get_network_speed(context: RunContext) -> str:
    return subprocess.run(["powershell", "Get-NetAdapter | Select Name, LinkSpeed"], capture_output=True, text=True).stdout[:500]

@function_tool()
async def control_network_adapter(context: RunContext, action: str) -> str:
    cmd = f'netsh interface set interface "Wi-Fi" {action}'
    subprocess.run(cmd, shell=True)
    return f"Adapter {action}d"

@function_tool()
async def remote_shutdown(context: RunContext, ip_address: str, username: str = "", password: str = "") -> str:
    os.system(f"shutdown /s /m \\\\{ip_address} /t 30")
    return "Shutdown command sent"

@function_tool()
async def remote_restart(context: RunContext, ip_address: str, username: str = "", password: str = "") -> str:
    os.system(f"shutdown /r /m \\\\{ip_address} /t 30")
    return "Restart command sent"
    
@function_tool()
async def execute_remote_command(context: RunContext, ip_address: str, command: str, username: str, password: str) -> str:
    return "Remote Execution via PowerShell not fully implemented in stub."

@function_tool()
async def remote_desktop_connect(context: RunContext, ip_address: str, username: str = "") -> str:
    subprocess.Popen(f"mstsc /v:{ip_address}", shell=True)
    return "RDP Launched"

@function_tool()
async def get_remote_device_info(context: RunContext, ip_address: str) -> str:
    try:
        host = socket.gethostbyaddr(ip_address)[0]
        return f"Host: {host}"
    except: return "Host not found"

@function_tool()
async def send_network_message(context: RunContext, ip_address: str, message: str) -> str:
    subprocess.run(f'msg /SERVER:{ip_address} * "{message}"', shell=True)
    return "Message sent"

@function_tool()
async def share_folder_on_network(context: RunContext, folder_path: str, share_name: str) -> str:
    subprocess.run(f'net share {share_name}="{folder_path}" /GRANT:Everyone,FULL', shell=True)
    return "Shared"

@function_tool()
async def remove_network_share(context: RunContext, share_name: str) -> str:
    subprocess.run(f'net share {share_name} /DELETE', shell=True)
    return "Share removed"

@function_tool()
async def list_network_shares(context: RunContext, ip_address: str = "") -> str:
    cmd = f"net view \\\\{ip_address}" if ip_address else "net share"
    return subprocess.run(cmd, shell=True, capture_output=True, text=True).stdout[:1000]

@function_tool()
async def access_network_share(context: RunContext, ip_address: str, share_name: str) -> str:
    os.startfile(f"\\\\{ip_address}\\{share_name}")
    return "Opened"

@function_tool()
async def map_network_drive(context: RunContext, drive_letter: str, network_path: str) -> str:
    subprocess.run(f"net use {drive_letter}: {network_path}", shell=True)
    return "Mapped"

@function_tool()
async def disconnect_network_drive(context: RunContext, drive_letter: str) -> str:
    subprocess.run(f"net use {drive_letter}: /delete", shell=True)
    return "Disconnected"

@function_tool()
async def get_device_services(context: RunContext, ip_address: str) -> str:
    return await port_scan(RunContext(), ip_address, "21,22,80,443,445,3389")

@function_tool()
async def remote_file_copy(context: RunContext, local_file: str, remote_ip: str, remote_path: str) -> str:
    import shutil
    shutil.copy2(local_file, f"\\\\{remote_ip}\\{remote_path}")
    return "Copied"

@function_tool()
async def control_jio_router(context: RunContext, action: str, password: str = "") -> str:
    if action == "open_portal":
        webbrowser.open("http://192.168.29.1")
        return "Opened Jio Router Portal"
    return "Action not supported in stub"

@function_tool()
async def disconnect_device_from_internet(context: RunContext, ip_address: str, mac_address: str = "") -> str:
    return "⚠️ ADMIN REQUIRED. Scapy dependency needed."

@function_tool()
async def reconnect_device_to_internet(context: RunContext, ip_address: str, mac_address: str = "") -> str:
    return "⚠️ ADMIN REQUIRED. Scapy dependency needed."

@function_tool()
async def check_admin_privileges(context: RunContext) -> str:
    import ctypes
    return "Admin" if ctypes.windll.shell32.IsUserAnAdmin() else "Not Admin"

@function_tool()
async def enable_packet_forwarding(context: RunContext) -> str:
    return "Packet forwarding enable command sent (requires admin)"

@function_tool()
async def get_router_admin_page(context: RunContext) -> str:
    webbrowser.open("http://192.168.0.1")
    return "Opened"

@function_tool()
async def network_kill_switch(context: RunContext, action: str) -> str:
    return "Kill switch toggled (simulated)"

@function_tool()
async def monitor_network_traffic(context: RunContext, duration_seconds: int = 10) -> str:
    return await network_bandwidth_usage(RunContext())

@function_tool()
async def limit_device_bandwidth(context: RunContext, ip_address: str, limit_kbps: int) -> str:
    return "Bandwidth limited (requires QoS)"

@function_tool()
async def remove_bandwidth_limit(context: RunContext, ip_address: str) -> str:
    return "Limit removed"

@function_tool()
async def get_network_devices_with_names(context: RunContext) -> str:
    return await scan_network_devices(RunContext())
