package main

import (
	"crypto/tls"
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"time"

	_ "github.com/mattn/go-sqlite3"
)

type ClientOperData struct {
	CiscoIOSXEWirelessClientOper struct {
		CommonOperData []struct {
			ClientMac string `json:"client-mac"`
			ApName    string `json:"ap-name"`
			WlanID    int    `json:"wlan-id"`
		} `json:"common-oper-data"`
		Dot11OperData []struct {
			ClientMac string `json:"ms-mac-address"`
			VapSSID   string `json:"vap-ssid"`
		} `json:"dot11-oper-data"`
		DCInfo []struct {
			ClientMac    string `json:"client-mac"`
			DeviceType   string `json:"device-type"`
			DayZeroDC    string `json:"day-zero-dc"`
			DeviceName   string `json:"device-name"`
			DeviceVendor string `json:"device-vendor"`
		} `json:"dc-info"`
	} `json:"Cisco-IOS-XE-wireless-client-oper:client-oper-data"`
}

func isAndroidDevice(mac string) (bool, string) {
	macUpper := strings.ToUpper(mac)
	parts := strings.Split(macUpper, ":")
	if len(parts) < 3 {
		return false, ""
	}

	oui := strings.Join(parts[:3], ":")

	if manufacturer, exists := androidManufacturerOUIs[strings.ToLower(oui)]; exists {
		return true, manufacturer
	}

	return false, ""
}

func main() {
	showFlag := flag.Bool("show", false, "Show stored clients")
	androidOnlyFlag := flag.Bool("android", false, "Show only Android devices (use with -show)")
	ouiFilterFlag := flag.Bool("oui", false, "Filter by Android manufacturer OUI when fetching")
	apFilterFlag := flag.String("ap", "", "Filter by AP name (use with -show)")
	ssidFilterFlag := flag.String("ssid", "", "Filter by SSID (use with -show)")
	flag.Parse()

	db, err := initDB()
	if err != nil {
		log.Fatal("Failed to initialize database:", err)
	}
	defer db.Close()

	if *showFlag {
		if err := showClients(db, *androidOnlyFlag, *apFilterFlag, *ssidFilterFlag); err != nil {
			log.Fatal("Failed to show clients:", err)
		}
		return
	}

	data, err := fetchClientData()
	if err != nil {
		log.Fatal("Failed to fetch client data:", err)
	}

	var clientsToSave []map[string]string
	if *ouiFilterFlag {
		clientsToSave = filterByOUI(data)
		fmt.Printf("Filtered %d Android devices by OUI\n", len(clientsToSave))
	} else {
		clientsToSave = getAllClients(data)
	}

	if err := saveToDatabase(db, clientsToSave); err != nil {
		log.Fatal("Failed to save to database:", err)
	}

	fmt.Printf("Successfully saved %d clients to database\n", len(clientsToSave))
}

func initDB() (*sql.DB, error) {
	db, err := sql.Open("sqlite3", "./clients.db")
	if err != nil {
		return nil, err
	}

	createTableSQL := `CREATE TABLE IF NOT EXISTS android_clients (
		id INTEGER PRIMARY KEY AUTOINCREMENT,
		client_mac TEXT NOT NULL,
		ap_name TEXT,
		ssid TEXT,
		device_type TEXT,
		day_zero_dc TEXT,
		device_name TEXT,
		device_vendor TEXT,
		timestamp DATETIME DEFAULT CURRENT_TIMESTAMP
	);`

	_, err = db.Exec(createTableSQL)
	if err != nil {
		return nil, err
	}

	return db, nil
}

func fetchClientData() (*ClientOperData, error) {
	tr := &http.Transport{
		TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
	}
	client := &http.Client{
		Transport: tr,
		Timeout:   60 * time.Second,
	}

	req, err := http.NewRequest("GET", "https://10.0.3.4/restconf/data/Cisco-IOS-XE-wireless-client-oper:client-oper-data", nil)
	if err != nil {
		return nil, err
	}

	req.SetBasicAuth("othman.suseno", "Pnd77net!")
	req.Header.Set("Accept", "application/yang-data+json")

	fmt.Println("Fetching data from Cisco API...")
	resp, err := client.Do(req)
	if err != nil {
		return nil, fmt.Errorf("request failed: %w", err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return nil, fmt.Errorf("API returned status code: %d", resp.StatusCode)
	}

	fmt.Println("Reading response...")
	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read response: %w", err)
	}

	fmt.Printf("Received %d bytes, parsing JSON...\n", len(body))

	var data ClientOperData
	if err := json.Unmarshal(body, &data); err != nil {
		return nil, fmt.Errorf("failed to parse JSON: %w", err)
	}

	return &data, nil
}

func getAllClients(data *ClientOperData) []map[string]string {
	var allClients []map[string]string

	dcInfoMap := make(map[string]map[string]string)
	for _, dc := range data.CiscoIOSXEWirelessClientOper.DCInfo {
		dcInfoMap[dc.ClientMac] = map[string]string{
			"device_type":   dc.DeviceType,
			"day_zero_dc":   dc.DayZeroDC,
			"device_name":   dc.DeviceName,
			"device_vendor": dc.DeviceVendor,
		}
	}

	dot11InfoMap := make(map[string]string)
	for _, dot11 := range data.CiscoIOSXEWirelessClientOper.Dot11OperData {
		dot11InfoMap[dot11.ClientMac] = dot11.VapSSID
	}

	fmt.Printf("Found %d clients with device classification info\n", len(dcInfoMap))
	fmt.Printf("Found %d clients with SSID info\n", len(dot11InfoMap))
	fmt.Printf("Found %d total clients in common-oper-data\n", len(data.CiscoIOSXEWirelessClientOper.CommonOperData))

	for i, client := range data.CiscoIOSXEWirelessClientOper.CommonOperData {
		dcInfo, hasDCInfo := dcInfoMap[client.ClientMac]
		ssid, hasSSID := dot11InfoMap[client.ClientMac]

		if i < 3 {
			fmt.Printf("Sample client %d: MAC=%s, AP=%s, SSID=%s\n",
				i+1, client.ClientMac, client.ApName, ssid)
		}

		clientData := map[string]string{
			"client_mac":    client.ClientMac,
			"ap_name":       client.ApName,
			"ssid":          "",
			"device_type":   "",
			"day_zero_dc":   "",
			"device_name":   "",
			"device_vendor": "",
		}

		if hasSSID {
			clientData["ssid"] = ssid
		}

		if hasDCInfo {
			clientData["device_type"] = dcInfo["device_type"]
			clientData["day_zero_dc"] = dcInfo["day_zero_dc"]
			clientData["device_name"] = dcInfo["device_name"]
			clientData["device_vendor"] = dcInfo["device_vendor"]
		}

		allClients = append(allClients, clientData)
	}

	return allClients
}

func filterByOUI(data *ClientOperData) []map[string]string {
	var androidClients []map[string]string

	dcInfoMap := make(map[string]map[string]string)
	for _, dc := range data.CiscoIOSXEWirelessClientOper.DCInfo {
		dcInfoMap[dc.ClientMac] = map[string]string{
			"device_type":   dc.DeviceType,
			"day_zero_dc":   dc.DayZeroDC,
			"device_name":   dc.DeviceName,
			"device_vendor": dc.DeviceVendor,
		}
	}

	dot11InfoMap := make(map[string]string)
	for _, dot11 := range data.CiscoIOSXEWirelessClientOper.Dot11OperData {
		dot11InfoMap[dot11.ClientMac] = dot11.VapSSID
	}

	androidCount := 0
	for _, client := range data.CiscoIOSXEWirelessClientOper.CommonOperData {
		isAndroid, manufacturer := isAndroidDevice(client.ClientMac)

		if isAndroid {
			dcInfo, hasDCInfo := dcInfoMap[client.ClientMac]
			ssid, hasSSID := dot11InfoMap[client.ClientMac]

			clientData := map[string]string{
				"client_mac":    client.ClientMac,
				"ap_name":       client.ApName,
				"ssid":          "",
				"device_type":   "",
				"day_zero_dc":   "",
				"device_name":   "",
				"device_vendor": manufacturer,
			}

			if hasSSID {
				clientData["ssid"] = ssid
			}

			if hasDCInfo {
				clientData["device_type"] = dcInfo["device_type"]
				clientData["day_zero_dc"] = dcInfo["day_zero_dc"]
				clientData["device_name"] = dcInfo["device_name"]
				if dcInfo["device_vendor"] != "" {
					clientData["device_vendor"] = dcInfo["device_vendor"]
				}
			}

			androidClients = append(androidClients, clientData)
			androidCount++

			if androidCount <= 5 {
				fmt.Printf("Found Android device: MAC=%s, AP=%s, SSID=%s, Manufacturer=%s\n",
					client.ClientMac, client.ApName, clientData["ssid"], manufacturer)
			}
		}
	}

	return androidClients
}

func saveToDatabase(db *sql.DB, clients []map[string]string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}
	defer tx.Rollback()

	stmt, err := tx.Prepare(`INSERT INTO android_clients 
		(client_mac, ap_name, ssid, device_type, day_zero_dc, device_name, device_vendor) 
		VALUES (?, ?, ?, ?, ?, ?, ?)`)
	if err != nil {
		return err
	}
	defer stmt.Close()

	for _, client := range clients {
		_, err = stmt.Exec(
			client["client_mac"],
			client["ap_name"],
			client["ssid"],
			client["device_type"],
			client["day_zero_dc"],
			client["device_name"],
			client["device_vendor"],
		)
		if err != nil {
			return err
		}
	}

	return tx.Commit()
}

func showClients(db *sql.DB, androidOnly bool, apFilter string, ssidFilter string) error {
	query := "SELECT client_mac, ap_name, ssid, device_type, day_zero_dc, device_name, device_vendor FROM android_clients"
	var args []interface{}
	var conditions []string

	if apFilter != "" {
		conditions = append(conditions, "ap_name LIKE ?")
		args = append(args, "%"+apFilter+"%")
	}

	if ssidFilter != "" {
		conditions = append(conditions, "ssid LIKE ?")
		args = append(args, "%"+ssidFilter+"%")
	}

	if len(conditions) > 0 {
		query += " WHERE " + strings.Join(conditions, " AND ")
	}

	rows, err := db.Query(query, args...)
	if err != nil {
		return err
	}
	defer rows.Close()

	fmt.Println("\n=== Stored Clients ===")
	fmt.Printf("%-20s %-20s %-20s %-15s %-15s %-30s %-20s\n",
		"MAC Address", "AP Name", "SSID", "Device Type", "Day Zero DC", "Device Name", "Device Vendor")
	fmt.Println(strings.Repeat("-", 140))

	count := 0
	androidCount := 0
	for rows.Next() {
		var mac, apName, ssid, deviceType, dayZeroDC, deviceName, deviceVendor string
		if err := rows.Scan(&mac, &apName, &ssid, &deviceType, &dayZeroDC, &deviceName, &deviceVendor); err != nil {
			return err
		}

		isAndroid, manufacturer := isAndroidDevice(mac)

		if androidOnly && !isAndroid {
			continue
		}

		if isAndroid {
			androidCount++
			if deviceVendor == "" {
				deviceVendor = manufacturer
			}
		}

		fmt.Printf("%-20s %-20s %-20s %-15s %-15s %-30s %-20s\n",
			mac, apName, ssid, deviceType, dayZeroDC, deviceName, deviceVendor)
		count++
	}

	fmt.Println(strings.Repeat("-", 140))
	if androidOnly {
		fmt.Printf("Total Android devices: %d\n", count)
	} else {
		fmt.Printf("Total clients: %d (Android: %d)\n", count, androidCount)
	}

	return rows.Err()
}