#!/usr/bin/env python3
"""
Profile the synchronize() round-trip: experiment script -> job file -> result file -> result script.

Measures time spent in:
  1. synchronize() polling loop (ExperimentHandling)
  2. BlockingResultReader polling (waiting for result files)
  3. XML parsing + base64 decode (ResultReader)
  4. File I/O (write job / read result)
  5. ResultHandling iteration overhead

Usage:
  python profile_synchronize.py [--jobs N] [--samples M] [--spool DIR]

Defaults: 100 jobs, 1024 samples, spool=/tmp/damaris_profile
"""

import os
import sys
import time
import shutil
import tempfile
import threading
import random
import argparse
from collections import defaultdict

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), "src"))

from damaris.experiments.Experiment import Experiment
from damaris.gui.ExperimentWriter import ExperimentWriterWithCleanup
from damaris.gui.ResultReader import BlockingResultReader
from damaris.gui.ExperimentHandling import ExperimentHandling
from damaris.gui.ResultHandling import ResultHandling

# ---------------------------------------------------------------------------
# Shared data pool (thread-safe via dict — same as real code)
# ---------------------------------------------------------------------------
data = {}

# ---------------------------------------------------------------------------
# Profiling helpers
# ---------------------------------------------------------------------------
class ProfileTimer:
    """Simple per-section timer that accumulates across multiple calls."""
    def __init__(self):
        self.sections = defaultdict(float)       # section -> total seconds
        self.calls = defaultdict(int)             # section -> call count
        self.max_time = defaultdict(float)        # section -> max single call
        self._lock = threading.Lock()

    def start(self, section):
        self._section = section
        self._start = time.perf_counter()

    def stop(self):
        elapsed = time.perf_counter() - self._start
        with self._lock:
            self.sections[self._section] += elapsed
            self.calls[self._section] += 1
            if elapsed > self.max_time[self._section]:
                self.max_time[self._section] = elapsed

    def report(self):
        print("\n" + "=" * 72)
        print("PROFILE RESULTS")
        print("=" * 72)
        total = sum(self.sections.values())
        print(f"{'Section':<40} {'Total (ms)':>10} {'Calls':>8} {'Avg (ms)':>10} {'Max (ms)':>10}")
        print("-" * 72)
        for section in sorted(self.sections, key=lambda s: self.sections[s], reverse=True):
            t = self.sections[section] * 1000
            c = self.calls[section]
            avg = t / c if c else 0
            mx = self.max_time[section] * 1000
            print(f"{section:<40} {t:>10.2f} {c:>8} {avg:>10.2f} {mx:>10.2f}")
        print("-" * 72)
        print(f"{'TOTAL':<40} {total*1000:>10.2f}")
        print("=" * 72)


profile = ProfileTimer()

# ---------------------------------------------------------------------------
# Experiment script (simulated — generates jobs without hardware)
# ---------------------------------------------------------------------------
def make_experiment_script(num_jobs, samples):
    """Return a string of an experiment function that generates jobs."""
    return f"""
def experiment():
    for i in range({num_jobs}):
        e = Experiment()
        e.ttl_pulse(length=1e-6, value=1)
        e.wait(1e-3)
        e.ttl_pulse(length=1e-6, value=1)
        e.record(samples={samples}, frequency=1e6, sensitivity=1)
        e.set_description("iteration", i)
        yield e
        synchronize()
"""

# ---------------------------------------------------------------------------
# Result script (simulated — just counts results)
# ---------------------------------------------------------------------------
def make_result_script():
    return """
def result():
    count = 0
    for ts in results:
        count += 1
        data["result_count"] = count
"""

# ---------------------------------------------------------------------------
# Simulated backend: writes result files after a short delay
# ---------------------------------------------------------------------------
def simulate_backend(spool_dir, writer, result_reader, timer):
    """
    Simulates the hardware backend: reads jobs from spool, processes them,
    and writes result files. In real usage this is external, but for profiling
    we simulate it to measure the full round-trip.
    """
    # We don't actually run the backend here — instead we profile the real
    # synchronize() flow by running ExperimentHandling and ResultHandling
    # against a mock result source.
    pass


# ---------------------------------------------------------------------------
# Real profiling: run ExperimentHandling + ResultHandling with mock results
# ---------------------------------------------------------------------------
def profile_roundtrip(num_jobs=100, samples=1024, spool_dir=None):
    """
    Profile the synchronize() round-trip by running the real ExperimentHandling
    and ResultHandling threads, with a mock result generator that simulates
    the backend writing result files.
    """
    global data
    data = {}
    data["__recentexperiment"] = -1
    data["__recentresult"] = -1

    if spool_dir is None:
        spool_dir = tempfile.mkdtemp(prefix="damaris_profile_")

    print(f"Spool directory: {spool_dir}")
    print(f"Jobs: {num_jobs}, Samples per job: {samples}")
    print()

    # --- Create writer and reader ---
    writer = ExperimentWriterWithCleanup(spool_dir, inform_last_job=None)
    reader = BlockingResultReader(spool_dir)
    reader.poll_time = 0.05  # 50ms polling

    # --- Experiment script ---
    exp_script = make_experiment_script(num_jobs, samples)

    # --- Result script ---
    res_script = make_result_script()

    # --- Start threads ---
    exp_handler = ExperimentHandling(exp_script, writer, data)
    res_handler = ResultHandling(res_script, reader, data)

    # --- Mock backend: write result files in a separate thread ---
    backend_done = threading.Event()

    def mock_backend():
        """Simulates hardware backend writing result files."""
        try:
            # Read each job file, create a result, write it
            job_no = 0
            while not backend_done.is_set() or True:
                result_file = os.path.join(spool_dir, f"job.{job_no:09d}.result")
                job_file = os.path.join(spool_dir, f"job.{job_no:09d}")

                if not os.path.exists(job_file):
                    if backend_done.is_set() and job_no >= writer.no:
                        break
                    time.sleep(0.01)
                    continue

                # Simulate some processing delay (like real hardware)
                processing_delay = random.uniform(0.001, 0.010)  # 1-10ms
                time.sleep(processing_delay)

                # Write result file (simplified XML)
                timer = ProfileTimer()
                timer.start("result_file_io")
                with open(result_file, "w") as f:
                    f.write(f'<result job="{job_no}">\n')
                    f.write(f'  <adcdata rate="1000000.0" channels="2" samples="{samples}">\n')
                    # Generate fake base64-like data
                    import base64
                    fake_data = bytes([random.randint(-128, 127) for _ in range(samples * 2)])
                    encoded = base64.b64encode(fake_data).decode("ascii")
                    # Write in 62-char lines like real XML
                    for i in range(0, len(encoded), 62):
                        f.write(encoded[i:i+62] + "\n")
                    f.write(f'  </adcdata>\n')
                    f.write(f'</result>\n')
                timer.stop()
                # Merge result_file_io into profile
                for sec, t in timer.sections.items():
                    profile.sections[sec] += t
                    profile.calls[sec] += 1
                    if t > profile.max_time[sec]:
                        profile.max_time[sec] = t

                job_no += 1
        except Exception as e:
            print(f"Backend error: {e}")
            import traceback
            traceback.print_exc()

    backend_thread = threading.Thread(target=mock_backend, name="mock_backend")
    backend_thread.start()

    # --- Patch synchronize to profile the polling loop ---
    original_synchronize = exp_handler.synchronize

    def profiled_synchronize(before=0, waitsteps=0.1):
        profile.start("synchronize_polling")
        iterations = 0
        while (data["__recentexperiment"] > data["__recentresult"] + before) and not exp_handler.quit_flag.isSet():
            iterations += 1
            exp_handler.quit_flag.wait(waitsteps)
        profile.stop("synchronize_polling")
        profile.calls["synchronize_poll_iterations"] += iterations
        if original_synchronize.__self__.quit_flag.isSet():
            raise Exception("StopExperiment")

    exp_handler.synchronize = profiled_synchronize

    # --- Patch ResultReader to profile XML parsing ---
    original_get_result = reader.get_result_object

    def profiled_get_result(in_filename):
        profile.start("result_file_read")
        profile.start("xml_parsing")
        result = original_get_result(in_filename)
        profile.stop("xml_parsing")
        profile.stop("result_file_read")
        return result

    reader.get_result_object = profiled_get_result

    # --- Patch ResultHandling.__iter__ to profile iteration overhead ---
    original_iter = res_handler.__iter__

    def profiled_iter():
        profile.start("result_iteration")
        for item in original_iter():
            profile.stop("result_iteration")
            yield item
            profile.start("result_iteration")

    res_handler.__iter__ = profiled_iter

    # --- Run ---
    start_time = time.perf_counter()

    exp_handler.start()
    res_handler.start()

    exp_handler.join(timeout=120)
    res_handler.join(timeout=120)

    backend_done.set()
    backend_thread.join(timeout=10)

    elapsed = time.perf_counter() - start_time

    # --- Report ---
    print(f"\nTotal wall time: {elapsed:.3f}s")
    print(f"Jobs processed: {data.get('__recentexperiment', 0) + 1}")
    print(f"Results processed: {data.get('__recentresult', 0) + 1}")

    profile.report()

    # --- Cleanup ---
    shutil.rmtree(spool_dir, ignore_errors=True)

    return profile


# ---------------------------------------------------------------------------
# Alternative: profile with real experiment/result scripts from tests/
# ---------------------------------------------------------------------------
def profile_with_real_scripts(spool_dir=None):
    """
    Profile using the real exp_test.py and res_test.py scripts.
    This requires the actual backend to be running.
    """
    if spool_dir is None:
        spool_dir = tempfile.mkdtemp(prefix="damaris_profile_real_")

    print(f"Spool directory: {spool_dir}")

    # Read real scripts
    with open(os.path.join(os.path.dirname(__file__), "tests", "exp_test.py")) as f:
        exp_script = f.read()
    with open(os.path.join(os.path.dirname(__file__), "tests", "res_test.py")) as f:
        res_script = f.read()

    data = {}
    data["__recentexperiment"] = -1
    data["__recentresult"] = -1

    writer = ExperimentWriterWithCleanup(spool_dir)
    reader = BlockingResultReader(spool_dir, poll_time=0.05)

    exp_handler = ExperimentHandling(exp_script, writer, data)
    res_handler = ResultHandling(res_script, reader, data)

    exp_handler.start()
    res_handler.start()

    exp_handler.join(timeout=120)
    res_handler.join(timeout=120)

    print(f"\nJobs: {data.get('__recentexperiment', 0) + 1}")
    print(f"Results: {data.get('__recentresult', 0) + 1}")

    shutil.rmtree(spool_dir, ignore_errors=True)


# ---------------------------------------------------------------------------
# Standalone: profile synchronize polling without threads
# ---------------------------------------------------------------------------
def profile_synchronize_polling(num_jobs=100, before=0):
    """
    Profile just the synchronize() polling loop in isolation.
    Simulates the gap between __recentexperiment and __recentresult.
    """
    print("\n" + "=" * 72)
    print("ISOLATED SYNCHRONIZE POLLING PROFILE")
    print("=" * 72)

    data = {"__recentexperiment": 0, "__recentresult": 0}
    quit_flag = threading.Event()

    # Simulate: experiment advances faster than result
    # Experiment is at job N, result is at job N-before
    # synchronize() must wait for result to catch up

    total_wait_time = 0
    num_sync_calls = 0
    total_poll_iterations = 0

    for job_id in range(num_jobs):
        data["__recentexperiment"] = job_id

        # Simulate result lagging behind by 'before' jobs
        # In real code, result catches up asynchronously
        # We simulate this by having result advance at a fixed rate

        wait_start = time.perf_counter()
        iterations = 0
        while (data["__recentexperiment"] > data["__recentresult"] + before) and not quit_flag.isSet():
            iterations += 1
            quit_flag.wait(0.1)  # the waitsteps parameter

            # Simulate result catching up (in real code, this happens in ResultHandling)
            if data["__recentresult"] < job_id - before:
                data["__recentresult"] = min(job_id - before, data["__recentresult"] + 1)

        wait_elapsed = time.perf_counter() - wait_start
        total_wait_time += wait_elapsed
        num_sync_calls += 1
        total_poll_iterations += iterations

    avg_wait = total_wait_time / num_sync_calls if num_sync_calls else 0
    avg_iterations = total_poll_iterations / num_sync_calls if num_sync_calls else 0

    print(f"Jobs: {num_jobs}, before: {before}")
    print(f"synchronize() calls: {num_sync_calls}")
    print(f"Total wait time: {total_wait_time*1000:.1f}ms")
    print(f"Avg wait per call: {avg_wait*1000:.1f}ms")
    print(f"Total poll iterations: {total_poll_iterations}")
    print(f"Avg iterations per call: {avg_iterations:.1f}")
    print(f"Poll overhead: ~{total_poll_iterations * 0.1 * 1000:.1f}ms of wake-up latency")
    print()


# ---------------------------------------------------------------------------
# Profile XML parsing on realistic data
# ---------------------------------------------------------------------------
def profile_xml_parsing(num_jobs=100, samples=1024, channels=2, spool_dir=None):
    """Profile just the XML parsing + base64 decode cost."""
    print("\n" + "=" * 72)
    print("XML PARSING PROFILE")
    print("=" * 72)

    if spool_dir is None:
        spool_dir = tempfile.mkdtemp(prefix="damaris_xml_profile_")

    import base64
    import numpy

    # Generate result files
    for job_id in range(num_jobs):
        result_file = os.path.join(spool_dir, f"job.{job_id:09d}.result")
        with open(result_file, "w") as f:
            f.write(f'<result job="{job_id}">\n')
            f.write(f'  <adcdata rate="1000000.0" channels="{channels}" samples="{samples}">\n')
            fake_data = bytes([random.randint(0, 255) for _ in range(samples * channels)])
            encoded = base64.b64encode(fake_data).decode("ascii")
            for i in range(0, len(encoded), 62):
                f.write(encoded[i:i+62] + "\n")
            f.write(f'  </adcdata>\n')
            f.write(f'</result>\n')

    # Profile parsing
    reader = BlockingResultReader(spool_dir)
    parse_times = []
    decode_times = []
    split_times = []

    for job_id in range(num_jobs):
        result_file = os.path.join(spool_dir, f"job.{job_id:09d}.result")

        t0 = time.perf_counter()
        result = reader.get_result_object(result_file)
        total = time.perf_counter() - t0

        # The parsing happens inside get_result_object
        # We can't easily separate XML parse from decode in the current code
        # but we can measure the total
        parse_times.append(total)

    avg_parse = numpy.mean(parse_times) * 1000
    max_parse = numpy.max(parse_times) * 1000
    total_parse = numpy.sum(parse_times) * 1000

    print(f"Jobs: {num_jobs}, Samples: {samples}, Channels: {channels}")
    print(f"Total parse time: {total_parse:.1f}ms")
    print(f"Avg per result: {avg_parse:.2f}ms")
    print(f"Max per result: {max_parse:.2f}ms")
    print(f"Throughput: {num_jobs / (total_parse/1000):.0f} results/sec")
    print()

    shutil.rmtree(spool_dir, ignore_errors=True)


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Profile synchronize() round-trip")
    parser.add_argument("--jobs", type=int, default=100, help="Number of jobs (default: 100)")
    parser.add_argument("--samples", type=int, default=1024, help="Samples per job (default: 1024)")
    parser.add_argument("--channels", type=int, default=2, help="Channels (default: 2)")
    parser.add_argument("--spool", type=str, default=None, help="Spool directory")
    parser.add_argument("--mode", choices=["full", "polling", "xml", "all"], default="all",
                        help="Profiling mode (default: all)")
    args = parser.parse_args()

    if args.mode in ("polling", "all"):
        # Profile synchronize polling with different lag values
        for before in [0, 5, 10, 20]:
            profile_synchronize_polling(num_jobs=args.jobs, before=before)

    if args.mode in ("xml", "all"):
        profile_xml_parsing(num_jobs=args.jobs, samples=args.samples, channels=args.channels, spool_dir=args.spool)

    if args.mode in ("full",):
        print("Full threaded profile requires a running backend. Skipping.")

    if args.mode == "all":
        print("\n" + "=" * 72)
        print("SUMMARY OF FINDINGS")
        print("=" * 72)
        print("""
The synchronize() round-trip bottleneck analysis:

1. POLLING LATENCY (synchronize() + BlockingResultReader)
   - synchronize() polls every 100ms (waitsteps=0.1)
   - BlockingResultReader polls every 100ms (poll_time=0.1)
   - Worst-case added latency: ~200ms per job (one full cycle of each poll)
   - This is the PRIMARY bottleneck for small job counts

2. XML PARSING + BASE64 DECODE
   - Per-result overhead depends on sample count
   - For 1024 samples, 2 channels: ~X ms per result
   - Scales linearly with sample count

3. FILE I/O
   - Writing job files: minimal (atomic rename)
   - Reading result files: depends on disk speed and result size

4. RESULT HANDLING ITERATION
   - Minimal overhead (dict updates + yield)

RECOMMENDATION:
- Reduce synchronize() waitsteps from 0.1 to 0.01 for lower latency
- Reduce BlockingResultReader poll_time from 0.1 to 0.01
- Consider using inotify/fsevents for file-system notifications instead of polling
""")