package main

import (
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"sort"
	"strings"

	"ringer/extract"
	"ringer/pefile"
	"ringer/runtime"
)

// sysIoctl is the JSON form of an extracted IOCTL code.
type sysIoctl struct {
	Code       uint32 `json:"code"`
	DeviceType uint16 `json:"device_type"`
	DeviceName string `json:"device_type_name"`
	Function   uint16 `json:"function"`
	Method     uint8  `json:"method"`
	MethodName string `json:"method_name"`
	Access     uint8  `json:"access"`
	AccessName string `json:"access_name"`
	Confidence string `json:"confidence"`
	Source     string `json:"source"`
	RVA        uint32 `json:"rva"`
}

// sysString is the JSON form of an extracted string.
type sysString struct {
	Value    string `json:"value"`
	Encoding string `json:"encoding"`
	Offset   uint32 `json:"offset"`
	Category string `json:"category"`
}

// systemDriver is the per-driver result of a system check.
type systemDriver struct {
	Name        string      `json:"name"`
	Path        string      `json:"path"`
	ImageBase   string      `json:"image_base"`
	ImageSize   uint32      `json:"image_size"`
	Ioctls      []sysIoctl  `json:"ioctls"`
	RingBuffers []sysString `json:"ring_buffers"`
	Devices     []sysString `json:"devices"`
}

// systemCheckReport is the full system-check output.
type systemCheckReport struct {
	TotalDrivers int            `json:"total_drivers"`
	Drivers      []systemDriver `json:"drivers"`
}

// runSystemCheck audits every loaded kernel module, writes a JSON report to
// outPath, and prints a summary to stdout.
func runSystemCheck(outPath string, minString int) error {
	mods, err := runtime.Modules()
	if err != nil {
		return fmt.Errorf("enumerate loaded modules: %w", err)
	}

	report := systemCheckReport{TotalDrivers: len(mods), Drivers: make([]systemDriver, 0, len(mods))}
	for _, m := range mods {
		d := systemDriver{
			Name:      filepath.Base(m.Path),
			Path:      m.Path,
			ImageBase: fmt.Sprintf("0x%x", m.ImageBase),
			ImageSize: m.ImageSize,
		}
		if f, err := pefile.Open(resolveModulePath(m.Path)); err == nil {
			for _, c := range extract.ExtractIoctls(f) {
				d.Ioctls = append(d.Ioctls, toSysIoctl(c))
			}
			for _, s := range extract.ExtractStrings(f, minString) {
				ss := toSysString(s)
				switch s.Category {
				case "section", "ringbuffer":
					d.RingBuffers = append(d.RingBuffers, ss)
				case "device", "symlink":
					d.Devices = append(d.Devices, ss)
				}
			}
		}
		report.Drivers = append(report.Drivers, d)
	}

	data, err := json.MarshalIndent(report, "", "  ")
	if err != nil {
		return fmt.Errorf("marshal report: %w", err)
	}
	if err := os.WriteFile(outPath, data, 0o644); err != nil {
		return fmt.Errorf("write report %q: %w", outPath, err)
	}

	printSystemSummary(report)
	return nil
}

// printSystemSummary prints a concise summary of the system check to stdout.
func printSystemSummary(r systemCheckReport) {
	totalIoctls, totalRB, totalDev, interesting := 0, 0, 0, 0
	for _, d := range r.Drivers {
		totalIoctls += len(d.Ioctls)
		totalRB += len(d.RingBuffers)
		totalDev += len(d.Devices)
		if len(d.Ioctls) > 0 || len(d.RingBuffers) > 0 || len(d.Devices) > 0 {
			interesting++
		}
	}

	fmt.Printf("System check: %d drivers loaded\n", r.TotalDrivers)
	fmt.Printf("================================\n")
	fmt.Printf("Total IOCTLs:          %d\n", totalIoctls)
	fmt.Printf("Total ring buffers:    %d\n", totalRB)
	fmt.Printf("Total device names:    %d\n", totalDev)
	fmt.Printf("Drivers with findings: %d\n", interesting)

	type row struct {
		name string
		io   int
		rb   int
		dev  int
	}
	var rows []row
	for _, d := range r.Drivers {
		if len(d.Ioctls) == 0 && len(d.RingBuffers) == 0 && len(d.Devices) == 0 {
			continue
		}
		rows = append(rows, row{d.Name, len(d.Ioctls), len(d.RingBuffers), len(d.Devices)})
	}
	sort.Slice(rows, func(i, j int) bool {
		if rows[i].io != rows[j].io {
			return rows[i].io > rows[j].io
		}
		return rows[i].name < rows[j].name
	})

	if len(rows) > 0 {
		fmt.Printf("\n%-40s %8s %8s %8s\n", "Driver", "IOCTLs", "RingBuf", "Devices")
		fmt.Printf("%-40s %8s %8s %8s\n", strings.Repeat("-", 40), "------", "------", "------")
		for _, r := range rows {
			fmt.Printf("%-40s %8d %8d %8d\n", r.name, r.io, r.rb, r.dev)
		}
	}
}

// toSysIoctl converts an extracted IOCTL to its JSON form.
func toSysIoctl(c extract.IoctlCode) sysIoctl {
	dt := extract.DeviceTypeName(c.DeviceType)
	if dt == "" {
		dt = fmt.Sprintf("0x%04X", c.DeviceType)
	}
	return sysIoctl{
		Code:       c.Code,
		DeviceType: c.DeviceType,
		DeviceName: dt,
		Function:   c.Function,
		Method:     c.Method,
		MethodName: extract.MethodName(c.Method),
		Access:     c.Access,
		AccessName: extract.AccessName(c.Access),
		Confidence: c.Confidence,
		Source:     c.Source,
		RVA:        c.RVA,
	}
}

// toSysString converts an extracted string to its JSON form.
func toSysString(s extract.StringInfo) sysString {
	return sysString{
		Value:    s.Value,
		Encoding: s.Encoding,
		Offset:   s.Offset,
		Category: s.Category,
	}
}

// resolveModulePath converts a kernel module path (e.g. \SystemRoot\...) to a
// filesystem path.
func resolveModulePath(p string) string {
	root := os.Getenv("SystemRoot")
	if root == "" {
		root = `C:\Windows`
	}
	if strings.HasPrefix(p, `\SystemRoot\`) {
		return root + strings.TrimPrefix(p, `\SystemRoot`)
	}
	if strings.HasPrefix(p, `SystemRoot\`) {
		return root + `\` + strings.TrimPrefix(p, `SystemRoot\`)
	}
	return p
}