"""NASDAQ TotalView-ITCH: Order Book Data Parsing — parse ITCH binary protocol into structured messages."""

import gzip
import os
import shutil
import struct
import warnings
from collections import Counter, defaultdict
from datetime import date, datetime
from pathlib import Path
from time import time

warnings.filterwarnings("ignore")

import polars as pl
from itch_message_specs import (
    FMT_DICT,
    MESSAGE_SPECS,
    NT_DICT,
    flush_to_parquet,
    parse_price4,
    parse_timestamp,
    print_message_formats,
)
from tqdm.auto import tqdm

from data import load_nasdaq_itch

SKIP_PARSING = False

# Data paths
# Canonical parsed messages location (both read and write target)
MESSAGE_DIR = load_nasdaq_itch(get_base_path=True)
MESSAGE_DIR.mkdir(parents=True, exist_ok=True)

# Raw ITCH binary input (from download script)
ITCH_RAW_DIR = MESSAGE_DIR.parent / "raw"

print(f"Raw ITCH data (input): {ITCH_RAW_DIR}")
print(f"Parsed messages (output): {MESSAGE_DIR}")
_raw_present = ITCH_RAW_DIR.exists() and next(ITCH_RAW_DIR.iterdir(), None) is not None
print(f"Raw data exists: {_raw_present}")
if not _raw_present:
    print("  (run NB00_itch_download.py first to fetch the ITCH binary)")

# Show message formats (loaded from utils/itch_message_specs.py)
print_message_formats()

# Create a sample Add Order message to demonstrate parsing
# Format: >HH6sQsI8sI (big-endian)

sample_add_order = struct.pack(
    ">HH6sQsI8sI",
    1234,  # stock_locate
    5678,  # tracking_number
    b"\x00\x00\x00\x00\x00\x01",  # timestamp (1 nanosecond)
    9876543210,  # order_reference_number
    b"B",  # buy_sell_indicator
    100,  # shares
    b"AAPL    ",  # stock (padded to 8 chars)
    1500000,  # price (150.0000 in price4 format)
)

print(f"Raw Add Order message ({len(sample_add_order)} bytes):")
print(f"  Hex: {sample_add_order.hex()}")

# Parse the binary data using our struct format
parsed = struct.unpack(FMT_DICT["A"], sample_add_order)
add_order = NT_DICT["A"]._make(parsed)

print("Parsed Add Order Message:")
print("-" * 40)
for field, value in add_order._asdict().items():
    if isinstance(value, bytes):
        value = value.decode("ascii").strip()
    print(f"  {field:25}: {value}")

# Apply conversions using helper functions from utils.itch_message_specs
ts_ns = parse_timestamp(add_order.timestamp)
price = parse_price4(add_order.price)

print(f"Timestamp: {ts_ns:,} nanoseconds = {ts_ns / 1e9:.9f} seconds after midnight")
print(f"Price: ${price:.4f}")

# Check what data is available locally
print("ITCH Data Pipeline Status:")
print("-" * 50)

# Step 1: Raw binary from download
raw_files = []
if ITCH_RAW_DIR.exists():
    raw_files = list(ITCH_RAW_DIR.glob("*.gz")) + list(ITCH_RAW_DIR.glob("*.bin"))
print(f"Raw binary files: {len(raw_files)}")
for f in raw_files:
    print(f"  {f.name} ({f.stat().st_size / 1e9:.2f} GB)")

# Step 2: Parsed messages (single uppercase letter = message type)
parsed_types = (
    [
        d
        for d in sorted(MESSAGE_DIR.iterdir())
        if d.is_dir() and len(d.name) == 1 and d.name.isupper()
    ]
    if MESSAGE_DIR.exists()
    else []
)
parsed_with_data = [d for d in parsed_types if list(d.glob("*.parquet"))]
print(f"Parsed message types: {len(parsed_with_data)}")
for msg_dir in parsed_with_data:
    name = MESSAGE_SPECS.get(msg_dir.name, {}).get("name", "Unknown")
    n_files = len(list(msg_dir.glob("*.parquet")))
    print(f"  {msg_dir.name} ({name}): {n_files} files")

# Validate: at minimum we need parsed data to continue
trade_dir = MESSAGE_DIR / "P"
assert trade_dir.exists() and list(trade_dir.glob("*.parquet")), (
    f"No parsed ITCH data at {MESSAGE_DIR}.\n"
    "To set up the data pipeline:\n"
    "  1. Download raw data:  uv run python data/equities/market/microstructure/nasdaq_itch_download.py\n"
    "  2. Parse (this notebook, Section 4) or use Rust parser (Section 6)\n"
    "  3. Parsed messages go to: data/equities/market/microstructure/nasdaq_itch/messages/"
)

trades = pl.read_parquet(trade_dir / "*.parquet")
print(f"\nLoaded {len(trades):,} trade messages")
print(f"Columns: {trades.columns}")

def _read_frame(f, pbar) -> tuple[str, bytes] | None:
    """Read one ITCH message frame: 2-byte length + 1-byte type + payload.

    Returns (msg_type, payload) on success, or None on EOF/truncation.
    """
    # 2-byte big-endian length prefix (message size including type byte)
    length_bytes = f.read(2)
    if len(length_bytes) < 2:
        return None
    pbar.update(2)

    msg_size = int.from_bytes(length_bytes, "big")

    # 1-byte message type
    msg_type_byte = f.read(1)
    if len(msg_type_byte) < 1:
        print(f"\nWarning: Truncated message at byte {f.tell()}, expected type byte")
        return None
    pbar.update(1)

    msg_type = msg_type_byte.decode("ascii")

    # Payload (msg_size includes type byte, so payload is msg_size - 1)
    payload = f.read(msg_size - 1)
    if len(payload) < msg_size - 1:
        print(f"\nWarning: Truncated payload for message type {msg_type}")
        return None
    pbar.update(msg_size - 1)

    return msg_type, payload

def _decode_message(msg_type: str, payload: bytes) -> dict | None:
    """Unpack binary payload into a dict, converting timestamps and strings.

    Returns parsed message dict, or None on struct error.
    """
    try:
        parsed = struct.unpack(FMT_DICT[msg_type], payload)
        msg = NT_DICT[msg_type]._make(parsed)._asdict()
    except struct.error:
        return None

    # Convert timestamp: nanoseconds since midnight
    if "timestamp" in msg:
        msg["timestamp"] = int.from_bytes(msg["timestamp"], "big")

    # Decode string fields
    for field, value in msg.items():
        if isinstance(value, bytes):
            msg[field] = value.decode("ascii").strip()

    return msg

def parse_itch_file(
    itch_file: Path,
    trading_day: date,
    output_dir: Path,
    max_buffered_messages: int = 10_000_000,
    max_messages: int | None = None,
) -> dict[str, int]:
    """Parse ITCH binary file and store messages as Parquet.

    Args:
        itch_file: Path to binary ITCH file (.bin, not .gz).
        trading_day: Trading date for timestamp construction.
        output_dir: Directory for Parquet output (one subdir per message type).
        max_buffered_messages: Flush threshold (total buffered messages).
        max_messages: Optional limit for testing.

    Returns:
        Dictionary with message type counts.
    """
    # Midnight timestamp for the trading day (ITCH timestamps are nanoseconds offset)
    base_ts = datetime(trading_day.year, trading_day.month, trading_day.day)
    file_counters: dict[str, int] = defaultdict(int)

    buffers = defaultdict(list)
    counts = Counter()
    file_size = itch_file.stat().st_size
    start_time = time()

    with (
        itch_file.open("rb") as f,
        tqdm(total=file_size, desc="Parsing ITCH", unit="B", unit_scale=True) as pbar,
    ):
        while True:
            if max_messages and sum(counts.values()) >= max_messages:
                print(f"\nLimit reached: {max_messages:,} messages")
                break

            frame = _read_frame(f, pbar)
            if frame is None:
                break
            msg_type, payload = frame
            counts[msg_type] += 1

            if msg_type not in FMT_DICT:
                continue

            msg = _decode_message(msg_type, payload)
            if msg is None:
                continue

            # Check for end of messages
            if msg_type == "S" and msg.get("event_code") == "C":
                print("\nEnd of Messages")
                flush_to_parquet(buffers, output_dir, base_ts, file_counters)
                break

            buffers[msg_type].append(msg)

            # Periodic flush
            if sum(len(v) for v in buffers.values()) >= max_buffered_messages:
                flush_to_parquet(buffers, output_dir, base_ts, file_counters)

    # Final flush
    if any(buffers.values()):
        flush_to_parquet(buffers, output_dir, base_ts, file_counters)

    elapsed = time() - start_time
    total = sum(counts.values())
    print(f"Parsed {total:,} messages in {elapsed:.1f}s ({total / elapsed:,.0f} msg/s)")

    return dict(counts)

# Locate ITCH data file (skip if SKIP_PARSING is set)
if SKIP_PARSING:
    print("SKIP_PARSING=True: skipping ITCH binary parsing (uses pre-parsed data)")
    itch_file = None
    counts = {}
else:
    # Clear any existing parsed data to avoid schema conflicts
    # (Different parser versions may produce different schemas)
    # Set ITCH_KEEP_EXISTING=1 to skip cleanup and use existing data
    clear_existing = os.environ.get("ITCH_KEEP_EXISTING", "0") != "1"
    if clear_existing and MESSAGE_DIR.exists() and list(MESSAGE_DIR.glob("*/part-*.parquet")):
        print(f"Clearing existing parsed data in {MESSAGE_DIR}")
        print("  (Set ITCH_KEEP_EXISTING=1 to keep existing data)")
        shutil.rmtree(MESSAGE_DIR)
        MESSAGE_DIR.mkdir(parents=True, exist_ok=True)

    # Find ITCH file (compressed or uncompressed)
    gz_files = list(ITCH_RAW_DIR.glob("*.gz")) if ITCH_RAW_DIR.exists() else []
    bin_files = list(ITCH_RAW_DIR.glob("*.bin")) if ITCH_RAW_DIR.exists() else []

    if not gz_files and not bin_files:
        raise FileNotFoundError(
            f"No raw ITCH binary found at {ITCH_RAW_DIR}.\n"
            "Download first:\n"
            "  uv run python data/equities/market/microstructure/nasdaq_itch_download.py"
        )

# Decompress if needed and extract trading date
if not SKIP_PARSING and (gz_files or bin_files):
    # Prefer uncompressed, otherwise decompress
    if bin_files:
        itch_file = bin_files[0]
        print(f"Found uncompressed: {itch_file.name}")
    else:
        gz_file = gz_files[0]
        itch_file = gz_file.with_suffix(".bin")

        if not itch_file.exists():
            print(f"Decompressing {gz_file.name}...")
            with gzip.open(gz_file, "rb") as f_in, open(itch_file, "wb") as f_out:
                shutil.copyfileobj(f_in, f_out)
            print(f"Created: {itch_file.name} ({itch_file.stat().st_size / 1e9:.1f} GB)")
        else:
            print(f"Found: {itch_file.name}")

    # Extract trading date from filename (format: MMDDYYYY.NASDAQ_ITCH50.bin)
    date_str = itch_file.stem.split(".")[0]
    if len(date_str) == 8 and date_str.isdigit():
        trading_day = date(int(date_str[4:8]), int(date_str[:2]), int(date_str[2:4]))
    else:
        trading_day = date(2020, 1, 30)  # Fallback

    print(f"Trading day: {trading_day}")

    # Parse ITCH file (full-day parse: ~22 min on the reference machine)
    if itch_file and itch_file.exists():
        counts = parse_itch_file(
            itch_file=itch_file,
            trading_day=trading_day,
            output_dir=MESSAGE_DIR,
            # max_messages=1_000_000,  # Remove this line for full parse
        )

        print("\nMessage counts:")
        for msg_type, count in sorted(counts.items(), key=lambda x: -x[1]):
            name = MESSAGE_SPECS.get(msg_type, {}).get("name", "Unknown")
            print(f"  {msg_type} ({name:25}): {count:>10,}")

# Message type distribution — use lazy scan to count without loading all data
print("Parsed Message Types:")
print("-" * 50)
for msg_dir in sorted(MESSAGE_DIR.iterdir()):
    if (
        msg_dir.is_dir()
        and len(msg_dir.name) == 1
        and msg_dir.name.isupper()
        and list(msg_dir.glob("*.parquet"))
    ):
        count = pl.scan_parquet(msg_dir / "*.parquet").select(pl.len()).collect().item()
        name = MESSAGE_SPECS.get(msg_dir.name, {}).get("name", "Unknown")
        print(f"  {msg_dir.name} ({name:25}): {count:>12,} messages")

# Schema compatibility check — verify we can read each message type
print("Schema Compatibility Check:")
print("-" * 50)
for msg_dir in sorted(MESSAGE_DIR.iterdir()):
    if (
        msg_dir.is_dir()
        and len(msg_dir.name) == 1
        and msg_dir.name.isupper()
        and list(msg_dir.glob("*.parquet"))
    ):
        try:
            sample = pl.scan_parquet(msg_dir / "*.parquet").head(5).collect()
            name = MESSAGE_SPECS.get(msg_dir.name, {}).get("name", "Unknown")
            print(f"  [OK] {msg_dir.name} ({name:25}): cols={list(sample.columns)[:4]}...")
        except Exception as e:
            print(f"  [FAIL] {msg_dir.name}: {e}")