#!/usr/bin/env python3
# Yêu cầu: Python3, rich, subprocess, perf

import subprocess, sys, os, time, sqlite3
from rich.console import Console
from rich.table import Table
from rich.panel import Panel
from rich.live import Live
from rich.prompt import Prompt

console = Console()

if len(sys.argv) < 3:
    console.print("Usage: python sqlite_profiler_ultimate_lock.py <db_file> <query1> [<query2> ...]", style="red")
    sys.exit(1)

db_file = sys.argv[1]
queries = sys.argv[2:]
timeline_data = []

# Step 1: Profile each query
console.print(f"[bold green]Profiling {len(queries)} queries on {db_file} ...[/bold green]")

for idx, query in enumerate(queries, 1):
    console.print(f"\n[cyan]Query {idx}:[/cyan] {query}")
    perf_log = f"perf_stat_{idx}.log"

    # Step 1a: Run perf stat for CPU/I/O/Cache
    cmd = [
        "perf", "stat", "-d",
        "-e", "cycles,cache-misses,block:rq_issue,block:rq_complete",
        "sqlite3", db_file, query
    ]
    subprocess.run(cmd, stderr=open(perf_log, "w"))

    # Step 1b: Measure duration + lock waits using busy_handler
    lock_wait_count = 0
    lock_wait_time = 0.0
    def busy_handler(count):
        nonlocal lock_wait_count, lock_wait_time, start_busy
        lock_wait_count += 1
        lock_wait_time += time.time() - start_busy
        return True  # retry

    conn = sqlite3.connect(db_file)
    conn.set_busy_handler(busy_handler)
    start_busy = time.time()
    start = time.time()
    try:
        conn.execute(query)
    except Exception as e:
        pass
    end = time.time()
    duration = end - start
    conn.close()

    # Step 1c: Parse perf log
    data = {}
    with open(perf_log) as f:
        for line in f:
            parts = line.strip().split()
            if len(parts) < 2: continue
            val, event = parts[0], parts[-1]
            try:
                data[event] = float(val.replace(',',''))
            except:
                continue
    cycles = data.get("cycles",0)
    cache_misses = data.get("cache-misses",0)
    io_ops = data.get("block:rq_issue",0) + data.get("block:rq_complete",0)

    timeline_data.append({
        "index": idx,
        "query": query,
        "duration": duration,
        "cycles": cycles,
        "cache_misses": cache_misses,
        "io_ops": io_ops,
        "lock_wait_count": lock_wait_count,
        "lock_wait_time": lock_wait_time
    })

# Step 2: Render timeline with stacked bars including Lock
def render_timeline(timeline, highlight_idx=None):
    term_width = os.get_terminal_size().columns - 30
    table = Table(title="Ultimate SQLite Profiler (CPU/I/O/Cache/Lock)", expand=True)
    table.add_column("Idx", justify="right")
    table.add_column("Timeline")
    table.add_column("Bottleneck")

    max_total = max(
        q['cycles'] + q['io_ops'] + q['cache_misses'] + q['lock_wait_time']*1e9
        for q in timeline
    ) or 1

    for q in timeline:
        # normalize metrics
        lock_scaled = q['lock_wait_time']*1e9  # scale lock time to match cycles
        cpu_len = int((q['cycles']/max_total)*term_width)
        io_len = int((q['io_ops']/max_total)*term_width)
        cache_len = int((q['cache_misses']/max_total)*term_width)
        lock_len = int((lock_scaled/max_total)*term_width)
        remaining_len = term_width - (cpu_len + io_len + cache_len + lock_len)
        if remaining_len < 0: remaining_len = 0

        bar = f"[red]{'█'*cpu_len}[/red][yellow]{'█'*io_len}[/yellow][blue]{'█'*cache_len}[/blue][magenta]{'█'*lock_len}[/magenta]{'-'*remaining_len}"

        if highlight_idx == q['index']:
            bar = f"[bold]{bar}[/bold]"

        # Bottleneck heuristic
        if q['lock_wait_time'] > 0.05:
            bottleneck = "[magenta]Lock-bound[/magenta]"
        elif q['io_ops'] > q['cycles'] and q['io_ops'] > q['cache_misses']:
            bottleneck = "[yellow]I/O-bound[/yellow]"
        elif q['cache_misses'] > q['cycles'] and q['cache_misses'] > q['io_ops']:
            bottleneck = "[blue]Memory-bound[/blue]"
        elif q['cycles'] > q['io_ops'] and q['cycles'] > q['cache_misses']:
            bottleneck = "[red]CPU-bound[/red]"
        else:
            bottleneck = "[green]Balanced[/green]"

        table.add_row(str(q['index']), bar, bottleneck)
    return table

def render_query_details(q):
    panel = Panel(f"""
[cyan]SQL:[/cyan] {q['query']}

[bold]Metrics:[/bold]
Cycles          : {q['cycles']:,}
Cache Misses    : {q['cache_misses']:,}
I/O Ops         : {q['io_ops']:,}
Duration        : {q['duration']:.3f} s
Lock Wait Count : {q['lock_wait_count']}
Lock Wait Time  : {q['lock_wait_time']:.3f} s
""", title=f"Query {q['index']} Details", border_style="magenta")
    return panel

# Step 3: Interactive session
highlight_idx = None
from rich.live import Live
with Live(render_timeline(timeline_data, highlight_idx), refresh_per_second=1) as live:
    while True:
        choice = Prompt.ask("\nEnter query index to view details (q to quit)")
        if choice.lower() == 'q':
            break
        try:
            idx = int(choice)
            highlight_idx = idx
            live.update(render_timeline(timeline_data, highlight_idx))
            q = next(filter(lambda x:x['index']==idx, timeline_data))
            console.print(render_query_details(q))
        except Exception as e:
            console.print(f"[red]Invalid index[/red]")

console.print("\n[green]Ultimate Profiler with Lock metric ended.[/green]")