"""Run the fine-tuned TinyBERT model on new addresses.

Standalone script — only needs `transformers`+`torch` (and optionally
`gazetteer_lookup.py` + a pincodes CSV for the administrative fields).
Unlike the t5/qwen models, this is token classification (BIO tagging), not
JSON generation — it always produces a well-formed field dict; there's no
"invalid JSON" failure mode to handle.

Usage:
    # Single address
    python inference_tinybert.py "FLAT NO.32, UTTARA TOWERS, MG ROAD GUWAHATI , Kamrup Unclassified AS 781029"

    # Batch from stdin (one address per line)
    cat addresses.txt | python inference_tinybert.py --stdin

    # Batch from a text file, output JSONL
    python inference_tinybert.py --file addresses.txt --out results.jsonl
"""
from __future__ import annotations

import argparse
import json
import sys

MAX_LENGTH = 160

ALL_FIELDS = (
    "houseNumber", "houseName", "poi", "street", "subsubLocality", "subLocality",
    "locality", "village", "subDistrict", "district", "city", "state", "pincode",
)

LABELS = ["O"] + [f"{prefix}-{field}" for field in ALL_FIELDS for prefix in ("B", "I")]
ID2LABEL = {i: label for i, label in enumerate(LABELS)}


def extract_fields(raw_address: str, offsets: list[tuple[int, int]], pred_label_ids: list[int]) -> dict:
    """Reconstruct the field dict from per-token BIO predictions by slicing the
    raw address text at each contiguous B-/I- run's char span — never
    `tokenizer.decode`, which would lose casing and introduce WordPiece-merge
    artifacts (e.g. "##" continuation glue)."""
    result = {f: None for f in ALL_FIELDS}
    current_field = None
    current_start = None
    current_end = None

    def flush():
        if current_field is not None and result[current_field] is None:
            result[current_field] = raw_address[current_start:current_end]

    for (start, end), label_id in zip(offsets, pred_label_ids):
        if start == end:
            continue
        label = ID2LABEL[label_id]
        if label == "O":
            flush()
            current_field = None
            continue
        prefix, field = label.split("-", 1)
        if prefix == "B" or field != current_field:
            flush()
            current_field, current_start, current_end = field, start, end
        else:
            current_end = end
    flush()
    return result


def load_model(model_id: str):
    from transformers import AutoModelForTokenClassification, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(model_id)
    model = AutoModelForTokenClassification.from_pretrained(model_id)
    model.eval()
    return model, tokenizer


def parse_address(model, tokenizer, raw_address: str) -> dict:
    import torch

    enc = tokenizer(
        raw_address, return_tensors="pt", return_offsets_mapping=True,
        truncation=True, max_length=MAX_LENGTH,
    )
    offsets = enc.pop("offset_mapping")[0].tolist()
    with torch.no_grad():
        out = model(**enc)
    pred_ids = out.logits[0].argmax(-1).tolist()
    result = extract_fields(raw_address, offsets, pred_ids)
    result["_raw_address"] = raw_address
    return result


def main():
    p = argparse.ArgumentParser(description="Parse Indian addresses using the fine-tuned TinyBERT model")
    p.add_argument("address", nargs="?", help="Single address to parse")
    p.add_argument("--stdin", action="store_true", help="Read addresses from stdin, one per line")
    p.add_argument("--file", help="Read addresses from a text file")
    p.add_argument("--out", help="Write JSONL output to file (default: stdout)")
    p.add_argument("--model", default=".", help="Model dir or HF repo id (default: current directory)")
    p.add_argument("--pincodes", default=None,
                    help="Path to a pincodes CSV (India Post format) — if given, adds "
                         "districtAdministrative/stateAdministrative/cityAdministrative "
                         "fields via deterministic pincode lookup (does not alter the "
                         "model's own district/state/city)")
    args = p.parse_args()

    import os

    model, tokenizer = load_model(args.model)

    gazetteer = None
    if args.pincodes and os.path.exists(args.pincodes):
        from gazetteer_lookup import load_gazetteer

        gazetteer = load_gazetteer(args.pincodes)

    if args.address:
        addresses = [args.address]
    elif args.stdin:
        addresses = [line.rstrip("\n") for line in sys.stdin if line.strip()]
    elif args.file:
        with open(args.file, encoding="utf-8") as f:
            addresses = [line.rstrip("\n") for line in f if line.strip()]
    else:
        p.print_help()
        return

    out_f = open(args.out, "w", encoding="utf-8") if args.out else None
    for addr in addresses:
        result = parse_address(model, tokenizer, addr)
        if gazetteer is not None:
            from gazetteer_lookup import add_administrative_fields

            raw_addr = result.pop("_raw_address")
            result = add_administrative_fields(result, gazetteer)
            result["_raw_address"] = raw_addr
        line = json.dumps(result, ensure_ascii=False)
        if out_f:
            out_f.write(line + "\n")
        else:
            print(line)

    if out_f:
        out_f.close()
        print(f"Wrote {len(addresses):,} results to {args.out}", file=sys.stderr)


if __name__ == "__main__":
    main()