package main

import (
    "bytes"
    "context"
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "net/http"
    "os"
    "os/exec"
    "path/filepath"
    "regexp"
    "strconv"
    "strings"
    "sync"
    "time"
    "unicode"

    "gopkg.in/yaml.v3"
)

// ==========================================
// 1. TYPES & CONFIGURATION
// ==========================================

// ConfigStructure représente le fichier config.yaml distant
type Config struct {
    ConfigVersion int `yaml:"config_version"`
    Runtime       struct {
        MaxTotalRuntime int `yaml:"max_total_runtime_s"`
        CommandTimeout  int `yaml:"command_timeout_s"`
        TempDir         string `yaml:"temp_dir"`
    } `yaml:"runtime"`
    Backend struct {
        URL string `yaml:"url"`
    } `yaml:"backend"`
    Collection struct {
        System struct {
            Enabled bool     `yaml:"enabled"`
            Items   []string `yaml:"items"`
        } `yaml:"system"`
        CPU struct {
            Enabled bool `yaml:"enabled"`
        } `yaml:"cpu"`
        RAM struct {
            Enabled bool `yaml:"enabled"`
        } `yaml:"ram"`
        Storage struct {
            Enabled bool `yaml:"enabled"`
        } `yaml:"storage"`
        Network struct {
            Enabled bool `yaml:"enabled"`
        } `yaml:"network"`
    } `yaml:"collection"`
    Benchmarks struct {
        Enabled bool              `yaml:"enabled"`
        Weights map[string]float64 `yaml:"weights"`
        CPU     struct {
            Enabled bool   `yaml:"enabled"`
            Tool    string `yaml:"tool"`
            Params  struct {
                CpuMaxPrime int `yaml:"cpu_max_prime"`
            } `yaml:"params"`
        } `yaml:"cpu_sysbench"`
        Memory struct {
            Enabled bool   `yaml:"enabled"`
            Tool    string `yaml:"tool"`
            Params  struct {
                TotalSize string `yaml:"total_size"`
            } `yaml:"params"`
        } `yaml:"memory_sysbench"`
        Disk struct {
            Enabled bool   `yaml:"enabled"`
            Tool    string `yaml:"tool"`
            Safety  struct {
                MaxRuntime int `yaml:"max_runtime_s"`
            } `yaml:"safety"`
        } `yaml:"disk_fio"`
        Network struct {
            Enabled bool   `yaml:"enabled"`
            Tool    string `yaml:"tool"`
            Server  string `yaml:"server"`
            Port    int    `yaml:"port"`
            Params  struct {
                Duration int `yaml:"duration_s"`
            } `yaml:"params"`
        } `yaml:"network_iperf3"`
    } `yaml:"benchmarks"`
}

// FinalPayload est la structure JSON envoyée au backend
type FinalPayload struct {
    DeviceIdentifier   string      `json:"device_identifier"`
    BenchClientVersion string      `json:"bench_client_version"`
    Hardware           Hardware    `json:"hardware"`
    Results            Results     `json:"results"`
    RawInfo            map[string]string `json:"raw_info,omitempty"`
}

type Hardware struct {
    CPU       CPUInfo      `json:"cpu"`
    RAM       RAMInfo      `json:"ram"`
    Storage   []DiskInfo   `json:"storage"`
    Network   []NetInfo    `json:"network"`
    System    SystemInfo   `json:"system"`
}

type Results struct {
    CPU        CPUResult    `json:"cpu"`
    Memory     MemResult    `json:"memory"`
    Disk       DiskResult   `json:"disk"`
    Network    NetResult    `json:"network"`
    GlobalScore float64     `json:"global_score"`
}

type SystemInfo struct {
    OS           string `json:"os"`
    Kernel       string `json:"kernel"`
    Arch         string `json:"architecture"`
    Hostname     string `json:"hostname"`
    Virtualization string `json:"virtualization"`
}

type CPUInfo struct {
    Model     string `json:"model"`
    Cores     int    `json:"cores"`
    Threads   int    `json:"threads"`
}

type RAMInfo struct {
    TotalMB int `json:"total_mb"`
}

type DiskInfo struct {
    Name   string `json:"name"`
    SizeGB string `json:"size_gb"`
    Type   string `json:"type"`
}

type NetInfo struct {
    Name      string `json:"name"`
    IP        string `json:"ip_address"`
    SpeedMbps string `json:"speed_mbps"`
}

type CPUResult struct {
    ScoreSingle float64 `json:"score_single"`
    ScoreMulti  float64 `json:"score_multi"`
}

type MemResult struct {
    Throughput float64 `json:"throughput_mib_s"`
    Score      float64 `json:"score"`
}

type DiskResult struct {
    ReadMBs  float64 `json:"read_mb_s"`
    WriteMBs float64 `json:"write_mb_s"`
    Score    float64 `json:"score"`
}

type NetResult struct {
    Upload   float64 `json:"upload_mbps"`
    Download float64 `json:"download_mbps"`
    Score    float64 `json:"score"`
}

// ==========================================
// 2. GLOBALS & UTILS
// ==========================================

var (
    cfg     Config
    version = "1.0.0"
    debug   = false
    dryRun  = false
)

// safeRun exécute une commande avec timeout et gestion d'erreur
func safeRun(ctx context.Context, name string, args ...string) (string, error) {
    if debug {
        fmt.Printf("[DEBUG] Exec: %s %v\n", name, args)
    }
    if dryRun {
        return "DRY RUN OUTPUT", nil
    }

    cmd := exec.CommandContext(ctx, name, args...)
    var stdout, stderr bytes.Buffer
    cmd.Stdout = &stdout
    cmd.Stderr = &stderr

    err := cmd.Run()
    if err != nil {
        return "", fmt.Errorf("%w: %s", err, stderr.String())
    }
    return stdout.String(), nil
}

func printProgress(step, total int, name, status string) {
    // Simple affichage ANSI compatible
    statusColor := "\033[32m" // Vert
    if status == "WARN" || status == "SKIPPED" {
        statusColor = "\033[33m" // Jaune
    } else if strings.Contains(status, "ERROR") {
        statusColor = "\033[31m" // Rouge
    }
    reset := "\033[0m"
    fmt.Printf("[%d/%d] %-20s %s%s%s\n", step, total, name, statusColor, status, reset)
}

// ==========================================
// 3. CONFIG LOADER
// ==========================================

func loadConfig(url string) error {
    var configData []byte
    var err error

    // Tentative HTTP
    if url != "" {
        client := &http.Client{Timeout: 10 * time.Second}
        resp, err := client.Get(url)
        if err == nil {
            defer resp.Body.Close()
            configData, err = io.ReadAll(resp.Body)
            if err == nil {
                fmt.Println("INFO: Configuration distante chargée.")
                goto parse
            }
        }
        fmt.Printf("WARN: Impossible de charger la config distante (%s), tentative fallback...\n", err)
    }

    // Fallback local (ou simulation pour cet exemple)
    // Dans le réel, on lirait /var/cache/bench-client/config.yaml
    fmt.Println("INFO: Utilisation de la configuration embarquée (fallback).")
    configData = []byte(defaultConfigYAML) // Utilisation d'une constante pour la démo

parse:
    return yaml.Unmarshal(configData, &cfg)
}

// ==========================================
// 4. COLLECTORS
// ==========================================

func collectSystemInfo(ctx context.Context) SystemInfo {
    info := SystemInfo{}
    
    // Hostname
    if h, err := os.Hostname(); err == nil {
        info.Hostname = h
    }

    // OS / Kernel
    // Lecture simple de /etc/os-release
    if out, err := safeRun(ctx, "sh", "-c", "cat /etc/os-release | grep PRETTY_NAME | cut -d= -f2"); err == nil {
        info.OS = strings.Trim(out, "\"\n")
    }
    if out, err := safeRun(ctx, "uname", "-r"); err == nil {
        info.Kernel = strings.TrimSpace(out)
    }
    if out, err := safeRun(ctx, "uname", "-m"); err == nil {
        info.Arch = strings.TrimSpace(out)
    }
    if out, err := safeRun(ctx, "systemd-detect-virt"); err == nil {
        info.Virtualization = strings.TrimSpace(out)
    }

    return info
}

func collectCPU(ctx context.Context) CPUInfo {
    cpu := CPUInfo{}
    if !cfg.Collection.CPU.Enabled {
        return cpu
    }

    out, err := safeRun(ctx, "lscpu")
    if err != nil {
        fmt.Println("WARN: lscpu failed")
        return cpu
    }

    // Parsing basique via regex
    reModel := regexp.MustCompile(`Model name:\s+(.*)`)
    reCores := regexp.MustCompile(`^CPU\(s\):\s+(\d+)`)
    reThreads := regexp.MustCompile(`Thread\(s\) per core:\s+(\d+)`)

    lines := strings.Split(out, "\n")
    for _, line := range lines {
        if m := reModel.FindStringSubmatch(line); m != nil {
            cpu.Model = strings.TrimSpace(m[1])
        }
        if m := reCores.FindStringSubmatch(line); m != nil {
            cpu.Cores, _ = strconv.Atoi(m[1])
        }
        if m := reThreads.FindStringSubmatch(line); m != nil {
            cpu.Threads, _ = strconv.Atoi(m[1])
        }
    }
    return cpu
}

func collectRAM(ctx context.Context) RAMInfo {
    ram := RAMInfo{}
    if !cfg.Collection.RAM.Enabled {
        return ram
    }

    out, err := safeRun(ctx, "free", "-m")
    if err != nil {
        return ram
    }
    
    // Parsing de 'free -m' (ligne Mem:)
    lines := strings.Split(out, "\n")
    if len(lines) > 1 {
        fields := strings.Fields(lines[1])
        if len(fields) > 1 {
            ram.TotalMB, _ = strconv.Atoi(fields[1])
        }
    }
    return ram
}

func collectStorage(ctx context.Context) []DiskInfo {
    disks := []DiskInfo{}
    if !cfg.Collection.Storage.Enabled {
        return disks
    }

    out, err := safeRun(ctx, "lsblk", "-d", "-n", "-o", "NAME,SIZE,ROTA")
    if err != nil {
        return disks
    }

    lines := strings.Split(out, "\n")
    for _, line := range lines {
        parts := strings.Fields(line)
        if len(parts) >= 3 {
            dType := "ssd"
            if parts[2] == "1" {
                dType = "hdd"
            }
            disks = append(disks, DiskInfo{
                Name:   parts[0],
                SizeGB: parts[1],
                Type:   dType,
            })
        }
    }
    return disks
}

func collectNetwork(ctx context.Context) []NetInfo {
    nets := []NetInfo{}
    if !cfg.Collection.Network.Enabled {
        return nets
    }

    // Récupérer les interfaces UP
    out, err := safeRun(ctx, "ip", "-j", "addr")
    if err != nil {
        return nets
    }

    // ip -j addr renvoie du JSON, facile à parser
    var ipData []map[string]interface{}
    json.Unmarshal([]byte(out), &ipData)

    for _, iface := range ipData {
        ifname, _ := iface["ifname"].(string)
        if ifname == "lo" {
            continue
        }
        
        operstate, _ := iface["operstate"].(string)
        if operstate != "UP" && operstate != "UNKNOWN" {
            continue
        }

        addrInfo, _ := iface["addr_info"].([]interface{})
        ipStr := ""
        if len(addrInfo) > 0 {
            firstAddr := addrInfo[0].(map[string]interface{})
            ipStr, _ = firstAddr["local"].(string)
        }

        nets = append(nets, NetInfo{
            Name: ifname,
            IP:   ipStr,
        })
    }
    return nets
}

// ==========================================
// 5. BENCHMARKS
// ==========================================

func runCPUBench(ctx context.Context) CPUResult {
    res := CPUResult{}
    if !cfg.Benchmarks.CPU.Enabled {
        return res
    }

    fmt.Println("   -> Single thread...")
    // sysbench cpu --threads=1 run
    args := []string{"cpu", "--threads=1", "--time=10", "run"}
    if out, err := safeRun(ctx, "sysbench", args...); err == nil {
        // Parse: events per second:
        re := regexp.MustCompile(`events per second:\s+([\d\.]+)`)
        if m := re.FindStringSubmatch(out); len(m) > 1 {
            val, _ := strconv.ParseFloat(m[1], 64)
            res.ScoreSingle = val
        }
    }

    fmt.Println("   -> Multi thread...")
    nproc := os.Getenv("GOMAXPROCS") // fallback
    if nproc == "" {
        if out, err := safeRun(ctx, "nproc"); err == nil {
            nproc = strings.TrimSpace(out)
        } else {
            nproc = "4" // default safe
        }
    }
    
    args = []string{"cpu", fmt.Sprintf("--threads=%s", nproc), "--time=10", "run"}
    if out, err := safeRun(ctx, "sysbench", args...); err == nil {
        re := regexp.MustCompile(`events per second:\s+([\d\.]+)`)
        if m := re.FindStringSubmatch(out); len(m) > 1 {
            val, _ := strconv.ParseFloat(m[1], 64)
            res.ScoreMulti = val
        }
    }
    return res
}

func runMemBench(ctx context.Context) MemResult {
    res := MemResult{}
    if !cfg.Benchmarks.Memory.Enabled {
        return res
    }
    
    args := []string{"memory", "--memory-block-size=1K", "--memory-total-size=1G", "run"}
    if out, err := safeRun(ctx, "sysbench", args...); err == nil {
        // Parse: MiB/sec
        re := regexp.MustCompile(`([\d\.]+]\s*MiB/sec`) // exemple: 1024.00 MiB/sec
        // Ou plus simple sur la ligne summary:
        re2 := regexp.MustCompile(`transferred:\s+([\d\.]+)\s+MiB`)
        
        if m := re2.FindStringSubmatch(out); len(m) > 1 {
            val, _ := strconv.ParseFloat(m[1], 64)
            res.Throughput = val
            res.Score = val // Score direct
        }
    }
    return res
}

func runDiskBench(ctx context.Context) DiskResult {
    res := DiskResult{}
    if !cfg.Benchmarks.Disk.Enabled {
        return res
    }

    // Utilisation d'un fichier temporaire
    tmpFile := filepath.Join(cfg.Runtime.TempDir, "bench.fio")
    defer os.Remove(tmpFile)

    args := []string{
        "--name=bench", "--ioengine=libaio", "--rw=randrw", "--bs=4k",
        "--direct=1", "--size=500M", "--numjobs=1", "--runtime=30",
        "--time_based", "--output-format=json", "--filename=" + tmpFile,
    }

    if out, err := safeRun(ctx, "fio", args...); err == nil {
        var fioJSON map[string]interface{}
        if err := json.Unmarshal([]byte(out), &fioJSON); err == nil {
            jobs := fioJSON["jobs"].([]interface{})
            if len(jobs) > 0 {
                job := jobs[0].(map[string]interface{})
                read := job["read"].(map[string]interface{})
                write := job["write"].(map[string]interface{})
                
                // bw_bytes are usually in bytes/sec
                rBw := read["bw"].(float64) / 1024 / 1024 // to MB/s
                wBw := write["bw"].(float64) / 1024 / 1024 
                
                res.ReadMBs = rBw
                res.WriteMBs = wBw
                res.Score = (rBw + wBw) / 2
            }
        }
    }
    return res
}

func runNetBench(ctx context.Context) NetResult {
    res := NetResult{}
    if !cfg.Benchmarks.Network.Enabled {
        return res
    }

    args := []string{
        "-c", cfg.Benchmarks.Network.Server,
        "-p", strconv.Itoa(cfg.Benchmarks.Network.Port),
        "-J", "-t", strconv.Itoa(cfg.Benchmarks.Network.Params.Duration),
    }

    if out, err := safeRun(ctx, "iperf3", args...); err == nil {
        var iperfJSON map[string]interface{}
        if err := json.Unmarshal([]byte(out), &iperfJSON); err == nil {
            // iperf3 -J structure check
            if sum, ok := iperfJSON["end"].(map[string]interface{}); ok {
                // Sender = Upload (local sending)
                if sumSent, ok := sum["sum_sent"].(map[string]interface{}); ok {
                    bits := sumSent["bits_per_second"].(float64)
                    res.Upload = bits / 1000 / 1000 // Mbps
                }
                // Receiver = Download (local receiving)
                if sumRecv, ok := sum["sum_received"].(map[string]interface{}); ok {
                    bits := sumRecv["bits_per_second"].(float64)
                    res.Download = bits / 1000 / 1000
                }
                res.Score = (res.Upload + res.Download) / 2
            }
        }
    }
    return res
}

// ==========================================
// 6. MAIN ORCHESTRATOR
// ==========================================

const defaultConfigYAML = `
config_version: 1
runtime:
  max_total_runtime_s: 480
  command_timeout_s: 30
  temp_dir: "/tmp"
backend:
  url: "http://10.0.0.50:8007/api/benchmark"
collection:
  system: { enabled: true }
  cpu: { enabled: true }
  ram: { enabled: true }
  storage: { enabled: true }
  network: { enabled: true }
benchmarks:
  enabled: true
  weights: { cpu: 0.4, memory: 0.2, disk: 0.2, network: 0.1, gpu: 0.1 }
  cpu_sysbench:
    enabled: true
    params: { cpu_max_prime: 20000 }
  memory_sysbench:
    enabled: true
    params: { total_size: "1G" }
  disk_fio:
    enabled: true
    safety: { max_runtime_s: 60 }
  network_iperf3:
    enabled: true
    server: "127.0.0.1" # Default pour test
    port: 5201
    params: { duration_s: 10 }
`

func calculateScore(h Hardware, r Results) float64 {
    // Normalisation arbitraire pour l'exemple
    // CPU: assume 2000 events/sec is 100 pts
    cpuScore := ((r.CPU.ScoreSingle + r.CPU.ScoreMulti) / 2) / 20
    memScore := r.Memory.Throughput / 100 // assume 1000 MB/s is 100 pts
    diskScore := r.Disk.Score / 10       // assume 500 MB/s is 50 pts
    netScore := r.Network.Score / 10      // assume 1000 Mbps is 100 pts

    w := cfg.Benchmarks.Weights
    return (cpuScore*w.CPU + memScore*w.Memory + diskScore*w.Disk + netScore*w.Network) * 100
}

func main() {
    // Flags
    configURL := flag.String("config-url", "", "URL de la config distante")
    flag.BoolVar(&debug, "debug", false, "Mode debug")
    flag.BoolVar(&dryRun, "dry-run", false, "Simulation d'exécution")
    flag.Parse()

    // 1. Chargement Config
    fmt.Println("[1/6] Configuration...")
    if err := loadConfig(*configURL); err != nil {
        fmt.Printf("FATAL: Erreur chargement config: %v\n", err)
        os.Exit(1)
    }
    printProgress(1, 6, "Configuration", "OK")

    // 2. Contexte global
    ctx, cancel := context.WithTimeout(context.Background(), time.Duration(cfg.Runtime.MaxTotalRuntime)*time.Second)
    defer cancel()

    // 3. Collecte Info
    fmt.Println("[2/6] Collecte informations système...")
    hw := Hardware{
        System:  collectSystemInfo(ctx),
        CPU:     collectCPU(ctx),
        RAM:     collectRAM(ctx),
        Storage: collectStorage(ctx),
        Network: collectNetwork(ctx),
    }
    printProgress(2, 6, "Collecte", "OK")

    // 4. Benchmarks
    fmt.Println("[3/6] Exécution Benchmarks...")
    results := Results{
        CPU:     runCPUBench(ctx),
        Memory:  runMemBench(ctx),
        Disk:    runDiskBench(ctx),
        Network: runNetBench(ctx),
    }
    printProgress(3, 6, "Benchmarks", "OK")

    // 5. Score
    fmt.Println("[4/6] Calcul Score...")
    results.GlobalScore = calculateScore(hw, results)
    printProgress(4, 6, "Score", fmt.Sprintf("%.2f", results.GlobalScore))

    // 6. Payload & Envoi
    fmt.Println("[5/6] Préparation Payload...")
    rawInfo := make(map[string]string)
    if debug {
        rawInfo["debug"] = "active"
    }

    payload := FinalPayload{
        DeviceIdentifier:   hw.System.Hostname,
        BenchClientVersion: version,
        Hardware:           hw,
        Results:            results,
        RawInfo:            rawInfo,
    }

    jsonData, err := json.MarshalIndent(payload, "", "  ")
    if err != nil {
        fmt.Printf("FATAL: JSON error: %v\n", err)
        os.Exit(1)
    }
    printProgress(5, 6, "JSON", "OK")

    fmt.Println("[6/6] Envoi Backend...")
    if !dryRun {
        token := os.Getenv("API_TOKEN")
        if token == "" {
            fmt.Println("WARN: API_TOKEN manquant. Envoi skipped.")
            // Pour debug, on affiche le JSON quand même
            fmt.Println(string(jsonData))
            return
        }

        req, _ := http.NewRequest("POST", cfg.Backend.URL, bytes.NewBuffer(jsonData))
        req.Header.Set("Content-Type", "application/json")
        req.Header.Set("Authorization", "Bearer "+token)

        client := &http.Client{Timeout: 15 * time.Second}
        resp, err := client.Do(req)
        if err != nil {
            printProgress(6, 6, "Envoi HTTP", "ERROR")
            fmt.Printf("Erreur: %v\n", err)
            // Sauvegarde locale
            os.WriteFile("/tmp/bench_payload_last.json", jsonData, 0644)
        } else {
            defer resp.Body.Close()
            if resp.StatusCode >= 200 && resp.StatusCode < 300 {
                printProgress(6, 6, "Envoi HTTP", "OK")
            } else {
                printProgress(6, 6, "Envoi HTTP", fmt.Sprintf("FAIL %d", resp.StatusCode))
            }
        }
    } else {
        fmt.Println("DRY RUN: Payload non envoyé.")
        fmt.Println(string(jsonData))
    }
}