#!/usr/bin/env python3
import json
import sys
import os
import re
from datetime import datetime, timedelta

# Browser window offset – applied exactly once per coordinate.
BROWSER_WINDOW_OFFSET = 91


def parse_timestamp(timestamp):
    """Convert ISO 8601 timestamp to a datetime object."""
    return datetime.strptime(timestamp, "%Y-%m-%dT%H:%M:%S.%fZ")


def _flush_global_labels(active_records, global_label_info):
    """
    Remove global label info for any object (group or single)
    that no longer has any active record.
    """
    active_bases = {k[0] for k in active_records.keys()}
    for base in list(global_label_info.keys()):
        if base not in active_bases:
            global_label_info.pop(base)


def assign_label(base_key, current_pos, global_label_info, label_counter):
    """
    Determine or update a label for a given base_key and its current position.
    """
    info = global_label_info.get(base_key)
    if not info:
        info = {"base_label": f"autolabel_{label_counter}", "last_pos": None, "next_suffix": 1}
        global_label_info[base_key] = info
        label_counter += 1
    if info["last_pos"] == current_pos:
        label = info["base_label"]
    else:
        if info["last_pos"] is None:
            label = info["base_label"]
        else:
            label = f"{info['base_label']}_{info['next_suffix']}"
            info["next_suffix"] += 1
        info["last_pos"] = current_pos
    return label, label_counter


def generate_relative_time_ias(json_data, output_ias="areas_of_interest.ias", offset=0):
    """
    Processes json_data to export IAS file and builds a mapping of labels to HTML.
    """
    print(f"Offset time applied: {offset} ms")
    active_records = {}
    global_label_info = {}
    ia_data = []
    mapping = {}
    # initial timestamp
    ip_entry = next((e for e in json_data if e.get("coordinates", {}).get("message") == "First 's' key press logged"), None)
    if not ip_entry:
        raise ValueError("No initial 's' key press logged.")
    ip_time = parse_timestamp(ip_entry["coordinates"]["timestamp"]) - timedelta(milliseconds=offset)
    def relative_time_ms(ts): return int((ts - ip_time).total_seconds() * 1000)
    id_counter = 1
    label_counter = 1
    for entry in json_data:
        ts = parse_timestamp(entry["timestamp"])
        current_ms = relative_time_ms(ts)
        coords = [c for c in entry.get("coordinates", {}).get("coordinates", []) if c.get("width") and c.get("height")]
        seen_keys = set()
        if not coords:
            for key, rec in list(active_records.items()): rec.update(end=current_ms); ia_data.append(rec); active_records.pop(key)
            _flush_global_labels(active_records, global_label_info)
            continue
        # group vs single
        if len(coords) > 1:
            sorted_coords = sorted(coords, key=lambda c: c["y"] + BROWSER_WINDOW_OFFSET)
            merged_html, ys, xs, rights, bottoms, last_y = "", [], [], [], [], None
            for c in sorted_coords:
                x = c["x"]; y = c["y"] + BROWSER_WINDOW_OFFSET
                if last_y is not None and y > last_y: merged_html += "\n"
                merged_html += c["html"]; last_y = y; ys.append(y); xs.append(x); rights.append(x + c.get("width",0)); bottoms.append(y + c.get("height",0))
            base_key, pos_key = merged_html, tuple(ys)
            geometry = {"x": min(xs), "y": min(ys), "right": max(rights), "bottom": max(bottoms)}
        else:
            c = coords[0]; x = c["x"]; y = c["y"] + BROWSER_WINDOW_OFFSET
            merged_html = c["html"]; base_key, pos_key = (merged_html, (c.get("width",0), c.get("height",0))), y
            geometry = {"x": x, "y": y, "right": x + c.get("width",0), "bottom": y + c.get("height",0)}
        full_key = (base_key, pos_key); seen_keys.add(full_key)
        label, label_counter = assign_label(base_key, pos_key, global_label_info, label_counter)
        if full_key not in active_records:
            active_records[full_key] = {"start": current_ms, "end": current_ms, "id": id_counter, "label": label, "html": merged_html, **geometry}
            mapping[label] = merged_html; id_counter += 1
        else:
            active_records[full_key]["end"] = current_ms
        for key, rec in list(active_records.items()):
            if key not in seen_keys:
                rec.update(end=current_ms); ia_data.append(rec); active_records.pop(key)
        _flush_global_labels(active_records, global_label_info)
    final_ms = relative_time_ms(parse_timestamp(json_data[-1]["timestamp"]))
    for rec in active_records.values(): rec.update(end=final_ms); ia_data.append(rec)
    # write IAS
    with open(output_ias, "w") as f:
        f.write("# IA\tstart_time\tend_time\tshape\tID\tx\ty\tright\tbottom\tlabel\n")
        for rec in ia_data:
            f.write(f"-{rec['start']}\t-{rec['end']}\tRECTANGLE\t{rec['id']}\t{rec['x']:.2f}\t{rec['y']:.2f}\t{rec['right']:.2f}\t{rec['bottom']:.2f}\t{rec['label']}\n")
    print(f"IAS file generated at '{output_ias}'")
    return mapping


def compute_offset_from_session_log(log_filename):
    """Compute offset = first 'S' key down minus recording start."""
    first_s = rec_start = None
    with open(log_filename) as f:
        for line in f:
            parts = line.strip().split("\t")
            if len(parts) < 4: continue
            ts, msg = parts[1], parts[3]
            if "KeyDown [S]" in msg and first_s is None: first_s = int(ts)
            if "Screen Recording: Starting component Screen Recording" in msg and rec_start is None: rec_start = int(ts)
            if first_s and rec_start: break
    if first_s is None or rec_start is None:
        raise ValueError("Required events not found in session log.")
    return first_s - rec_start


def main():
    if len(sys.argv) != 2:
        print("Usage: python div_and_session_log_to_IA_and_html.py <participant_id>")
        sys.exit(1)
    identifier = sys.argv[1]
    input_file = f"div_logs_{identifier}.json"
    output_ias = f"output_{identifier}.ias"
    output_html = f"html_{identifier}.json"
    # determine offset
    log_file = f"SessionLog_{identifier}.log"
    if os.path.exists(log_file):
        try:
            offset = compute_offset_from_session_log(log_file)
            print(f"Computed offset from session log '{log_file}': {offset} ms")
        except Exception as e:
            print(f"Error computing offset: {e}. Using 0 ms."); offset = 0
    else:
        offset = 0
        print(f"No session log '{log_file}'. Using offset: 0 ms")
    # load and run
    with open(input_file, encoding="utf-8") as f:
        data = json.load(f)
    mapping = generate_relative_time_ias(data, output_ias, offset)
    # write html mapping
    with open(output_html, "w", encoding="utf-8") as f:
        json.dump(mapping, f, indent=2, ensure_ascii=False)
    print(f"HTML mapping generated at '{output_html}'")

if __name__ == "__main__":
    main()

#
# This script expects a single argument: the participant identifier and the following files:
# - div_logs_{identifier}.json
# - SessionLog_{identifier}.log
#
# It will generate:
# - output_{identifier}.ias (Interest Area File)
# - html_{identifier}.json (Suggestion HTML content mapping)
#
# Usage: python div_and_session_log_to_IA_and_html.py <participant_id>
#