"""Exporta model.joblib a ONNX + inference.json (sin depender de services-count)."""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import joblib
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Alineado con services-count-facturas-medicas-ml v1.1.0 y pipeline Go.
FEATURE_COLUMNS = [
    "first_token_is_code",
    "first_token_is_service_code",
    "has_service_table_pattern",
    "has_header_keyword",
    "currency_at_row_end",
    "first_token_x_norm",
    "is_likely_continuation",
    "currency_count",
    "integer_count",
    "x_span_norm",
    "y_norm",
    "token_count",
    "special_char_ratio",
    "digit_ratio",
]

DEFAULT_OCR = {"dpi": 300, "lang": "spa"}
DEFAULT_DBSCAN = {"eps": 10, "min_samples": 1}


def parse_args() -> argparse.Namespace:
    root = Path(__file__).resolve().parents[1]
    parser = argparse.ArgumentParser(description="Exportar Random Forest joblib a ONNX")
    parser.add_argument(
        "--joblib",
        type=Path,
        required=True,
        help="Ruta al model.joblib (ej. models/v1.1.0/model.joblib)",
    )
    parser.add_argument(
        "--out",
        type=Path,
        default=root / "models" / "latest",
        help="Directorio de salida para model.onnx e inference.json",
    )
    parser.add_argument(
        "--metadata",
        type=Path,
        default=None,
        help="model_metadata.json opcional para pipeline_version y params OCR/DBSCAN",
    )
    return parser.parse_args()


def load_metadata(path: Path | None) -> dict:
    if path is None or not path.exists():
        return {}
    return json.loads(path.read_text(encoding="utf-8"))


def main() -> None:
    args = parse_args()
    if not args.joblib.exists():
        raise SystemExit(f"No se encontró: {args.joblib}")

    bundle = joblib.load(args.joblib)
    model = bundle["model"]
    threshold = float(bundle.get("threshold", 0.5))
    meta = load_metadata(args.metadata)

    bundle_columns = list(bundle.get("feature_columns") or FEATURE_COLUMNS)
    if bundle_columns != FEATURE_COLUMNS:
        raise SystemExit(
            f"feature_columns del joblib no coinciden con el pipeline Go.\n"
            f"  joblib: {bundle_columns}\n"
            f"  esperado: {FEATURE_COLUMNS}"
        )

    n_features = int(getattr(model, "n_features_in_", len(FEATURE_COLUMNS)))
    if n_features != len(FEATURE_COLUMNS):
        raise SystemExit(
            f"El model.joblib tiene {n_features} features pero el pipeline Go produce "
            f"{len(FEATURE_COLUMNS)}. Reentrena o usa un joblib compatible."
        )

    initial_type = [("float_input", FloatTensorType([None, len(FEATURE_COLUMNS)]))]
    options = {id(model): {"zipmap": False}}
    onnx_model = convert_sklearn(
        model,
        initial_types=initial_type,
        options=options,
        target_opset=12,
    )

    out_dir = args.out
    out_dir.mkdir(parents=True, exist_ok=True)
    onnx_path = out_dir / "model.onnx"
    meta_path = out_dir / "inference.json"

    with open(onnx_path, "wb") as f:
        f.write(onnx_model.SerializeToString())

    inference = {
        "threshold": threshold,
        "feature_columns": FEATURE_COLUMNS,
        "pipeline_version": meta.get("pipeline_version", bundle.get("pipeline_version", "unknown")),
        "ocr": meta.get("ocr", DEFAULT_OCR),
        "dbscan": meta.get("dbscan", DEFAULT_DBSCAN),
    }
    meta_path.write_text(json.dumps(inference, indent=2), encoding="utf-8")

    print(f"ONNX: {onnx_path}")
    print(f"Meta: {meta_path}")
    print(f"Umbral: {threshold}")
    print(f"Features: {len(FEATURE_COLUMNS)}")


if __name__ == "__main__":
    main()