import sys, time, json, os
sys.path.insert(0, "/data/knowledge/maritime")

print("=" * 60)
print("RAG Performance Analysis")
print("=" * 60)

# 1. Index load time
from search_engine import _load_index, _get_config, search, tokenize
cfg = _get_config()
print("\n--- Index Loading ---")
t0 = time.time()
idx = _load_index()
t1 = time.time()
print("Load time:", round(t1 - t0, 2), "s")
print("Chunks:", idx["total"])
print("Keywords:", len(idx["keywords"]))

# Index file sizes
for name in ["rag_index.json", "rag_index.keywords.json"]:
    path = os.path.join(cfg["base_dir"], name) if "base_dir" in cfg else "/data/knowledge/maritime/" + name
    if os.path.exists(path):
        size_mb = os.path.getsize(path) / 1024 / 1024
        print(f"  {name}: {size_mb:.1f} MB")

# 2. Search latency
print("\n--- Search Latency ---")
queries = [
    "船舶进出港报告制度",
    "海上交通安全法 水污染防治法 船舶管理",
    "船员适任证书申请条件",
    "2019年船员发展报告",
    "琼州海峡定线制",
    "上海海事大学船长及格率",
]
for q in queries:
    times = []
    for _ in range(3):
        t0 = time.time()
        results = search(q, top_k=5)
        elapsed = time.time() - t0
        times.append(elapsed)
    avg = sum(times) / len(times)
    top_score = results[0]["score"] if results else 0
    top_src = results[0]["source"][:40] if results else "N/A"
    print(f"  Q: {q[:30]:30s} | {avg*1000:6.0f}ms | score={top_score:.1f} | {top_src}")

# 3. Tokenization analysis
print("\n--- Tokenization ---")
for q in ["海上交通安全法", "船员适任证书"]:
    tokens = tokenize(q)
    print(f"  '{q}' -> {len(tokens)} tokens: {tokens[:10]}")

# 4. Keyword index density
print("\n--- Keyword Index Density ---")
kw = idx["keywords"]
lengths = [len(v) for v in kw.values()]
print(f"  Total keywords: {len(kw)}")
print(f"  Avg postings per keyword: {sum(lengths)/len(lengths):.1f}")
print(f"  Max postings: {max(lengths)} (keyword: {[k for k,v in kw.items() if len(v)==max(lengths)][0][:10]})")
print(f"  Keywords with 1 posting: {sum(1 for l in lengths if l == 1)}")
print(f"  Keywords with >1000 postings: {sum(1 for l in lengths if l > 1000)}")

# 5. Search quality issue: check what "海上交通安全法" retrieves
print("\n--- Search Quality Check ---")
r1 = search("海上交通安全法 船舶 航行", top_k=3)
for i, r in enumerate(r1):
    has_law = "海上交通安全法" in r["source"]
    print(f"  [{i+1}] {'[LAW]' if has_law else '[REF]'} score={r['score']:.1f} {r['source'][:60]}")

# 6. LLM call time (biggest bottleneck)
print("\n--- LLM Inference Time ---")
import urllib.request
body = json.dumps({
    "model": "qwen2.5-32b",
    "messages": [{"role": "user", "content": "Reply OK in one word"}],
    "max_tokens": 5,
}).encode("utf-8")
req = urllib.request.Request("http://127.0.0.1:8000/v1/chat/completions",
    data=body, headers={"Content-Type": "application/json"}, method="POST")
t0 = time.time()
with urllib.request.urlopen(req, timeout=30) as resp:
    data = json.loads(resp.read())
t1 = time.time()
print(f"  Simple query: {(t1-t0)*1000:.0f}ms")

# Longer prompt
body2 = json.dumps({
    "model": "qwen2.5-32b",
    "messages": [{"role": "user", "content": "用200字概括中国海事法律体系" + "x" * 2000}],
    "max_tokens": 500,
}).encode("utf-8")
req2 = urllib.request.Request("http://127.0.0.1:8000/v1/chat/completions",
    data=body2, headers={"Content-Type": "application/json"}, method="POST")
t0 = time.time()
with urllib.request.urlopen(req2, timeout=60) as resp:
    data2 = json.loads(resp.read())
t1 = time.time()
tokens_out = data2.get("usage", {}).get("completion_tokens", 0)
print(f"  Long prompt + 500 tokens: {(t1-t0)*1000:.0f}ms ({tokens_out} output tokens)")
if tokens_out > 0:
    print(f"  Tokens/sec: {tokens_out/(t1-t0):.1f}")

print("\n--- Summary ---")
print("  Bottleneck: LLM inference (10-20s) >> Index load (once) >> Search (<100ms)")