"""
Pass 1: collect unique ASINs from sampled_1m_reviews.jsonl
Pass 2: stream All_Amazon_Meta.json.gz (double-gzip), write matching records

Output:
  sampled_meta.jsonl   — meta records cho các ASIN trong 1M users
  meta_stats.json      — thống kê coverage
"""
import gzip, json, time
from pathlib import Path

REVIEWS_PATH = Path("/workspace/amazon/sampled_1m_reviews.jsonl")
META_GZ      = Path("/workspace/amazon/All_Amazon_Meta.json.gz")
OUT_META     = Path("/workspace/amazon/sampled_meta.jsonl")
OUT_STATS    = Path("/workspace/amazon/meta_stats.json")

# ── Pass 1: collect ASINs from reviews ────────────────────────────────────
print("Pass 1: collecting ASINs from sampled reviews...", flush=True)
t0 = time.time()
review_asins = set()
total_reviews = 0

with open(REVIEWS_PATH, "r", encoding="utf-8", errors="replace") as f:
    for line in f:
        line = line.strip()
        if not line:
            continue
        try:
            rec = json.loads(line)
        except:
            continue
        asin = rec.get("asin", "")
        if asin:
            review_asins.add(asin)
        total_reviews += 1
        if total_reviews % 3_000_000 == 0:
            print(f"  {total_reviews/1e6:.0f}M reviews, {len(review_asins):,} ASINs  ({time.time()-t0:.0f}s)", flush=True)

print(f"Pass 1 done in {time.time()-t0:.1f}s", flush=True)
print(f"  Reviews scanned : {total_reviews:,}", flush=True)
print(f"  Unique ASINs    : {len(review_asins):,}", flush=True)

# ── Pass 2: stream meta (double-gzip), write matching records ─────────────
print(f"\nPass 2: streaming meta (double-gzip)...", flush=True)
t1 = time.time()
meta_total = 0
meta_written = 0
found_asins = set()

# field quality counters
has_title   = 0
has_price   = 0
has_also_buy  = 0
has_also_view = 0
has_desc    = 0
has_brand   = 0
has_category = 0

with gzip.open(META_GZ, "rb") as outer, \
     gzip.open(outer, "rt", encoding="utf-8", errors="replace") as f, \
     open(OUT_META, "w", encoding="utf-8") as fout:

    for line in f:
        line = line.strip()
        if not line:
            continue
        meta_total += 1
        if meta_total % 2_000_000 == 0:
            print(f"  {meta_total/1e6:.0f}M meta rows, {meta_written:,} matched  ({time.time()-t1:.0f}s)", flush=True)

        try:
            rec = json.loads(line)
        except:
            continue

        asin = rec.get("asin", "")
        if asin not in review_asins:
            continue

        # clean up noisy fields before saving
        clean = {
            "asin"       : asin,
            "title"      : rec.get("title", ""),
            "brand"      : rec.get("brand", ""),
            "price"      : rec.get("price", ""),
            "category"   : rec.get("category") or [],
            "main_cat"   : rec.get("main_cat", ""),
            "description": rec.get("description") or [],
            "feature"    : rec.get("feature") or [],
            "also_buy"   : rec.get("also_buy") or [],
            "also_view"  : rec.get("also_view") or [],
            "rank"       : rec.get("rank", ""),
        }
        fout.write(json.dumps(clean, ensure_ascii=False) + "\n")
        found_asins.add(asin)
        meta_written += 1

        # field quality tracking
        if clean["title"].strip():       has_title  += 1
        if clean["price"] not in ("", None, "None"): has_price += 1
        if clean["also_buy"]:            has_also_buy  += 1
        if clean["also_view"]:           has_also_view += 1
        if clean["brand"].strip():       has_brand  += 1
        if clean["category"]:            has_category += 1
        desc = clean["description"]
        if isinstance(desc, list):
            desc = " ".join(desc)
        if str(desc).strip():            has_desc += 1

print(f"Pass 2 done in {time.time()-t1:.1f}s", flush=True)

# ── Coverage & stats ───────────────────────────────────────────────────────
missing_asins = review_asins - found_asins
coverage = len(found_asins) / len(review_asins) * 100 if review_asins else 0

print(f"\n{'='*50}")
print(f"RESULTS")
print(f"{'='*50}")
print(f"Review ASINs      : {len(review_asins):,}")
print(f"Meta matched      : {meta_written:,}  ({coverage:.1f}% coverage)")
print(f"Missing from meta : {len(missing_asins):,}  ({100-coverage:.1f}%)")
print(f"\nField quality (matched records):")
n = meta_written or 1
print(f"  title       : {has_title/n*100:.1f}%")
print(f"  brand       : {has_brand/n*100:.1f}%")
print(f"  price       : {has_price/n*100:.1f}%")
print(f"  description : {has_desc/n*100:.1f}%")
print(f"  category    : {has_category/n*100:.1f}%")
print(f"  also_buy    : {has_also_buy/n*100:.1f}%")
print(f"  also_view   : {has_also_view/n*100:.1f}%")

stats = {
    "review_asins": len(review_asins),
    "meta_matched": meta_written,
    "meta_coverage_pct": round(coverage, 2),
    "missing_asins": len(missing_asins),
    "field_quality": {
        "title":       round(has_title/n*100, 1),
        "brand":       round(has_brand/n*100, 1),
        "price":       round(has_price/n*100, 1),
        "description": round(has_desc/n*100, 1),
        "category":    round(has_category/n*100, 1),
        "also_buy":    round(has_also_buy/n*100, 1),
        "also_view":   round(has_also_view/n*100, 1),
    }
}
with open(OUT_STATS, "w") as f:
    json.dump(stats, f, indent=2)

print(f"\nOutput: {OUT_META}  ({OUT_META.stat().st_size/1e9:.1f} GB)")
print(f"Total time: {time.time()-t0:.1f}s")