# In src/fmatch/scale.py
# REPLACE the entire file content with this corrected version.

from __future__ import annotations
import codecs
import os
import pathlib
import tempfile
import threading
import time
import logging
from typing import Callable, Optional, Dict, Any
import duckdb
import pyarrow as pa
import pyarrow.parquet as pq
from fmatch.core.engine import match_pairs_arrow
from concurrent.futures import ThreadPoolExecutor, as_completed
import psutil

log = logging.getLogger(__name__)


# ───────────────────── helper utils ─────────────────────
def _quote_ident(s: str) -> str:
    return '"' + s.replace('"', '""') + '"'


def _normalise(col: str) -> str:
    """Lower‑case, strip whitespace, punctuation and any '(…)' suffix."""
    if not col:
        return ""
    base = col.split("(")[0]
    return "".join(ch for ch in base if ch.isalnum()).lower()


def _pick_block_cols(cfg, hdr_L, hdr_R):
    iterable = cfg if isinstance(cfg, list) else cfg.get("columns", [])
    norm_L = {_normalise(h): h for h in hdr_L}
    norm_R = {_normalise(h): h for h in hdr_R}

    for m in iterable:
        f_raw, c_raw = m.get("foreign"), m.get("crm")
        if f_raw and c_raw and f_raw in hdr_L and c_raw in hdr_R:
            return f_raw, c_raw
        if _normalise(f_raw) in norm_L and _normalise(c_raw) in norm_R:
            return norm_L[_normalise(f_raw)], norm_R[_normalise(c_raw)]

    for key in norm_L.keys() & norm_R.keys():
        return norm_L[key], norm_R[key]

    log.warning("Falling back to first column of each file.")
    return hdr_L[0], hdr_R[0]


def _ensure_utf8(path, enc):
    p = pathlib.Path(path)
    if not enc or enc.lower().replace("-", "") in {"utf8", "utf"}:
        return p
    fd, tmp = tempfile.mkstemp(suffix=p.suffix, prefix="utf8_")
    os.close(fd)
    with codecs.open(p, "r", enc, errors="replace") as fi, codecs.open(
        tmp, "w", "utf-8"
    ) as fo:
        for ln in fi:
            fo.write(ln)
    return pathlib.Path(tmp)


def _match_block_arrow(
    tbl_L: pa.Table,
    tbl_R: pa.Table,
    cfg: Dict[str, Any] | list,
    *,
    algorithm: str | None = None,
    min_score: int | None = None,
    s_id_col: Optional[str] = None,
    r_id_col: Optional[str] = None,
) -> pa.Table:
    if isinstance(cfg, list):
        cfg = {"columns": cfg}
    if algorithm:
        cfg["algorithm"] = algorithm
    if min_score is not None:
        cfg["threshold"] = min_score
        cfg["score_cutoff"] = min_score
        cfg["min_score"] = min_score
    return match_pairs_arrow(tbl_L, tbl_R, cfg, s_id_col=s_id_col, r_id_col=r_id_col)


def scale_match(
    path_left: str | pathlib.Path,
    path_right: str | pathlib.Path,
    mapping_cfg: dict | list,
    output_dir: str | pathlib.Path,
    *,
    algorithm: str | None = None,
    min_score: int | None = None,
    left_encoding: str = "utf-8",
    right_encoding: str = "utf-8",
    blocking_cols: Optional[tuple[str, str]] = None,
    progress_cb: Callable[[Any], None] = lambda *_: None,
    cancel_event: Optional[threading.Event] = None,
    duckdb_pragmas: Optional[Dict[str, str]] = None,
    block_prefix_len: int = 3,
    coarse_fallback: bool = True,
    s_id_col: Optional[str] = None,
    r_id_col: Optional[str] = None,
    **_ignored,
) -> pathlib.Path:
    """
    Out-of-core fuzzy matcher driven by a parallelized DuckDB streaming engine.
    """
    start = time.time()
    MAX_BLOCK_PAIRS = 2_000_000
    cancel_event = cancel_event or threading.Event()
    out_dir = pathlib.Path(output_dir).expanduser().resolve()
    out_dir.mkdir(parents=True, exist_ok=True)

    num_workers = max(1, (psutil.cpu_count(logical=False) or 4))
    log.info(f"Using {num_workers} parallel workers for block processing.")

    main_con = duckdb.connect()
    if duckdb_pragmas:
        for k, v in duckdb_pragmas.items():
            main_con.execute(f"PRAGMA {k}='{v}'")

    p_left = _ensure_utf8(path_left, left_encoding)
    p_right = _ensure_utf8(path_right, right_encoding)
    hdr_L = list(
        main_con.from_csv_auto(str(p_left), header=True, sample_size=1).columns
    )
    hdr_R = list(
        main_con.from_csv_auto(str(p_right), header=True, sample_size=1).columns
    )
    col_L, col_R = blocking_cols or _pick_block_cols(mapping_cfg, hdr_L, hdr_R)

    def _get_keys(con, prefix_len):
        log.info(f"Creating views and fetching keys with prefix length: {prefix_len}")

        def _expr(col, n):
            return f"substr(lower(regexp_replace({_quote_ident(col)},'[^A-Za-z0-9]','','g')),1,{n})"

        con.execute(
            f"CREATE OR REPLACE VIEW L AS SELECT *, {_expr(col_L, prefix_len)} AS block_key FROM read_csv_auto('{p_left.as_posix()}', header=True, ENCODING='utf-8');"
        )
        con.execute(
            f"CREATE OR REPLACE VIEW R AS SELECT *, {_expr(col_R, prefix_len)} AS block_key FROM read_csv_auto('{p_right.as_posix()}', header=True, ENCODING='utf-8');"
        )
        return [
            k[0]
            for k in con.execute(
                "SELECT DISTINCT L.block_key FROM L INNER JOIN R USING(block_key)"
            ).fetchall()
        ]

    keys = _get_keys(main_con, block_prefix_len)
    if not keys and coarse_fallback and block_prefix_len > 1:
        log.warning("No common keys; retrying with 1-char prefix.")
        block_prefix_len = 1
        keys = _get_keys(main_con, block_prefix_len)

    log.info(
        "Scale blocking on %s <-> %s  (keys=%d, prefix=%d)",
        col_L,
        col_R,
        len(keys),
        block_prefix_len,
    )
    main_con.close()

    total_items_to_process = len(keys)  # The total for ETA is the number of blocks
    processed_items = 0
    comparisons_done = 0
    master_schema: pa.Schema | None = None
    progress_lock = threading.Lock()

    def _progress(items_delta=0, comps_delta=0):
        nonlocal processed_items, comparisons_done
        with progress_lock:
            processed_items += items_delta
            comparisons_done += comps_delta
            # ++ FIX: Send a unified ("stats", payload) message tuple ++
            payload = {
                "items_delta": items_delta,
                "comparisons_delta": comps_delta,
                "total_items_for_eta_calc": total_items_to_process,
            }
            progress_cb(("stats", payload))
            # ++ END FIX ++

    def _process_single_block(key: str):
        thread_con = duckdb.connect()

        def _expr(col, n):
            return f"substr(lower(regexp_replace({_quote_ident(col)},'[^A-Za-z0-9]','','g')),1,{n})"

        thread_con.execute(
            f"CREATE VIEW L AS SELECT *, {_expr(col_L, block_prefix_len)} AS block_key FROM read_csv_auto('{p_left.as_posix()}', header=True, ENCODING='utf-8');"
        )
        thread_con.execute(
            f"CREATE VIEW R AS SELECT *, {_expr(col_R, block_prefix_len)} AS block_key FROM read_csv_auto('{p_right.as_posix()}', header=True, ENCODING='utf-8');"
        )

        if not key or not key.strip():
            thread_con.close()
            return 0, 0, None, 1, "Skipped empty key"  # Return 1 for items_delta

        tblL = thread_con.execute(
            f"SELECT * FROM L WHERE block_key='{key}'"
        ).fetch_arrow_table()
        tblR = thread_con.execute(
            f"SELECT * FROM R WHERE block_key='{key}'"
        ).fetch_arrow_table()

        bytes_in_block = tblL.nbytes + tblR.nbytes
        num_pairs = len(tblL) * len(tblR)

        matches = _match_block_arrow(
            tblL,
            tblR,
            mapping_cfg,
            algorithm=algorithm,
            min_score=min_score,
            s_id_col=s_id_col,
            r_id_col=r_id_col,
        )

        thread_con.close()
        return bytes_in_block, num_pairs, matches, 1, None  # Return 1 block processed

    _progress()

    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        futures = {executor.submit(_process_single_block, key): key for key in keys}

        for future in as_completed(futures):
            if cancel_event.is_set():
                break

            key = futures[future]
            try:
                (
                    bytes_processed,
                    comparisons_processed,
                    matches,
                    items_processed,
                    error_msg,
                ) = future.result()
                if error_msg:
                    log.warning(error_msg)
                _progress(
                    items_delta=items_processed, comps_delta=comparisons_processed
                )

                if matches is not None and matches.num_rows > 0:
                    with progress_lock:
                        if master_schema is None:
                            master_schema = matches.schema
                        else:
                            matches = matches.cast(master_schema)
                    pq.write_table(matches, out_dir / f"part-{key}.parquet")

            except Exception as exc:
                log.error(f"Block '{key}' generated an exception: {exc}", exc_info=True)

    log.info("scale_match finished in %.1fs  ->  %s", time.time() - start, out_dir)
    return out_dir
