package com.databuild.examples.basic_graph;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.JsonNode;

import java.util.ArrayList;
import java.util.List;
import java.io.File;
import java.util.Arrays;
import java.util.Collections;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.Random;

/**
 * Unified job that handles both configuration and execution via subcommands.
 */
public class UnifiedGenerateNumber {
    public static String BASE_PATH = "/tmp/databuild_test/examples/basic_graph/";

    public static void main(String[] args) {
        if (args.length < 1) {
            System.err.println("Usage: UnifiedGenerateNumber {config|exec} [args...]");
            System.exit(1);
        }

        String command = args[0];
        switch (command) {
            case "config":
                handleConfig(Arrays.copyOfRange(args, 1, args.length));
                break;
            case "exec":
                handleExec(Arrays.copyOfRange(args, 1, args.length));
                break;
            default:
                System.err.println("Unknown command: " + command);
                System.err.println("Usage: UnifiedGenerateNumber {config|exec} [args...]");
                System.exit(1);
        }
    }

    private static void handleConfig(String[] args) {
        if (args.length < 1) {
            System.err.println("Config mode requires partition ref");
            System.exit(1);
        }

        String partitionRef = args[0];
        
        try {
            ObjectMapper mapper = new ObjectMapper();
            
            // Create job configuration
            var config = mapper.createObjectNode();
            
            // Create outputs as PartitionRef objects
            var outputs = mapper.createArrayNode();
            var outputPartRef = mapper.createObjectNode();
            outputPartRef.put("str", partitionRef);
            outputs.add(outputPartRef);
            config.set("outputs", outputs);
            
            config.set("inputs", mapper.createArrayNode());
            config.set("args", mapper.createArrayNode().add("will").add("generate").add(partitionRef));
            config.set("env", mapper.createObjectNode().put("PARTITION_REF", partitionRef));
            
            var response = mapper.createObjectNode();
            response.set("configs", mapper.createArrayNode().add(config));
            
            System.out.println(mapper.writeValueAsString(response));
        } catch (Exception e) {
            System.err.println("Error creating config: " + e.getMessage());
            System.exit(1);
        }
    }

    private static void handleExec(String[] args) {
        if (args.length < 3) {
            System.err.println("Execute mode requires: will generate <partition_ref>");
            System.exit(1);
        }

        String partitionRef = args[2];
        
        try {
            // Generate a random number based on the partition ref
            MessageDigest md = MessageDigest.getInstance("SHA-256");
            byte[] hash = md.digest(partitionRef.getBytes(StandardCharsets.UTF_8));
            long seed = 0;
            for (int i = 0; i < 8; i++) {
                seed = (seed << 8) | (hash[i] & 0xFF);
            }
            
            Random random = new Random(seed);
            int randomNumber = random.nextInt(100) + 1;
            
            // Write to file - partitionRef is the full path
            File outputFile = new File(partitionRef);
            File outputDir = outputFile.getParentFile();
            if (outputDir != null) {
                outputDir.mkdirs();
            }
            try (FileWriter writer = new FileWriter(outputFile)) {
                writer.write(String.valueOf(randomNumber));
            }
            
            System.out.println("Generated number " + randomNumber + " for partition " + partitionRef);
            
        } catch (Exception e) {
            System.err.println("Error in execution: " + e.getMessage());
            System.exit(1);
        }
    }
}