from pathlib import Path
from datetime import datetime, timezone
import csv, orjson, os
from typing import Dict, Iterable

OUTPUT_DIR = Path(os.getenv("OUTPUT_DIR", "./out"))

def _today_dir() -> Path:
    return OUTPUT_DIR / datetime.now(timezone.utc).strftime("%Y-%m-%d")

def ensure_dirs():
    (OUTPUT_DIR / "latest").mkdir(parents=True, exist_ok=True)
    (_today_dir()).mkdir(parents=True, exist_ok=True)

def write_jsonl(broker: str, rows: Iterable[Dict]) -> int:
    ensure_dirs()
    count = 0
    path = _today_dir() / f"{broker}.jsonl"
    with path.open("wb") as f:
        for row in rows:
            f.write(orjson.dumps(row) + b"\n")
            count += 1
    return count

def write_csv_latest(broker: str, rows: Iterable[Dict]) -> None:
    ensure_dirs()
    rows = list(rows)
    if not rows:
        return
    # compute union of keys across rows
    keys = set()
    for r in rows:
        keys.update(r.keys())
    keys = sorted(keys)
    with (OUTPUT_DIR / "latest" / f"{broker}.csv").open("w", newline="", encoding="utf-8") as f:
        w = csv.DictWriter(f, fieldnames=keys)
        w.writeheader()
        for r in rows:
            w.writerow({k: r.get(k) for k in keys})

def compose_combined_latest():
    ensure_dirs()
    combined_path = OUTPUT_DIR / "latest" / "combined.jsonl"
    with combined_path.open("wb") as out:
        for p in _today_dir().glob("*.jsonl"):
            out.write(p.read_bytes())
