Tr_Code/tools/import_fault_codes.py

import argparse
import os
import sqlite3
from typing import Dict, Any, List, Optional


def find_col(columns: List[str], candidates: List[str]) -> Optional[str]:
    lowered = {c.lower(): c for c in columns}
    for name in candidates:
        if name.lower() in lowered:
            return lowered[name.lower()]
    return None


def main() -> None:
    parser = argparse.ArgumentParser(description="Import fault codes into local data.db from external SQLite table")
    parser.add_argument("--src-db", default="fault_codes.db", help="source sqlite file path")
    parser.add_argument("--src-table", default="fault_code_list", help="source table name")
    parser.add_argument("--manufacturer", default="woojin", help="manufacturer slug (woojin|rotem)")
    parser.add_argument("--default-section", default="fault", help="default section if missing (fault|tcms|emergency)")
    parser.add_argument("--truncate", action="store_true", help="delete existing faults for manufacturer before import")

    args = parser.parse_args()

    # destination DB (our app)
    try:
        from app import DATABASE_PATH  # type: ignore
    except Exception:
        DATABASE_PATH = os.path.join(os.path.dirname(__file__), "..", "data.db")

    dest_path = os.path.abspath(DATABASE_PATH)
    src_path = os.path.abspath(args.src_db)

    if not os.path.exists(src_path):
        raise SystemExit(f"Source DB not found: {src_path}")

    src = sqlite3.connect(src_path)
    src.row_factory = sqlite3.Row
    dst = sqlite3.connect(dest_path)
    dst.row_factory = sqlite3.Row
    dst.execute("PRAGMA foreign_keys = ON")

    cur = src.execute(f"SELECT * FROM {args.src_table}")
    rows = cur.fetchall()
    if not rows:
        print("No rows in source table.")
        return

    columns = rows[0].keys()
    code_col = find_col(columns, ["code", "Code", "fault_code", "고장코드", "CODE"]) or "code"
    title_col = find_col(columns, ["title", "name", "f.Name", "f_name", "고장명", "fName"]) or None
    details_col = find_col(columns, ["details", "detail", "고장상세", "desc", "설명"])  # optional
    action_col = find_col(columns, ["action", "조치", "응급조치", "response", "반응"])  # optional
    cond_col = find_col(columns, ["condition", "검지조건", "조건"])  # optional
    category_col = find_col(columns, ["category", "cat", "장치분류", "system", "type"])  # optional
    section_col = find_col(columns, ["section", "섹션", "tab", "kind"])  # optional

    print("Column mapping:")
    print("  code ->", code_col)
    print("  title ->", title_col)
    print("  details ->", details_col)
    print("  action ->", action_col)
    print("  condition ->", cond_col)
    print("  category ->", category_col)
    print("  section ->", section_col, f"(default={args.default_section})")

    # manufacturer ensure
    m = dst.execute("SELECT id FROM manufacturer WHERE slug=?", (args.manufacturer.lower(),)).fetchone()
    if not m:
        name = "우진" if args.manufacturer.lower() == "woojin" else "로템"
        dst.execute("INSERT INTO manufacturer (slug, name) VALUES (?, ?)", (args.manufacturer.lower(), name))
        m_id = dst.execute("SELECT last_insert_rowid() AS id").fetchone()[0]
    else:
        m_id = m["id"]

    if args.truncate:
        dst.execute("DELETE FROM fault WHERE manufacturer_id=?", (m_id,))
        dst.execute("DELETE FROM category WHERE manufacturer_id=?", (m_id,))

    # existing categories map
    cat_map: Dict[str, int] = {}
    for r in dst.execute("SELECT id, name FROM category WHERE manufacturer_id=?", (m_id,)):
        cat_map[r["name"]] = r["id"]

    def ensure_category(name: Optional[str]) -> Optional[int]:
        if not name:
            return None
        name = str(name).strip()
        if name == "":
            return None
        if name in cat_map:
            return cat_map[name]
        dst.execute("INSERT INTO category (manufacturer_id, name) VALUES (?, ?)", (m_id, name))
        cid = dst.execute("SELECT last_insert_rowid() AS id").fetchone()[0]
        cat_map[name] = cid
        return cid

    def get_value(r: sqlite3.Row, key: Optional[str]):
        if not key:
            return None
        for k in r.keys():
            if k.lower() == key.lower():
                return r[k]
        return None

    batch: List[tuple] = []
    for row in rows:
        code = str(get_value(row, code_col) or "").strip()
        title = str(get_value(row, title_col) or "").strip()
        details = str(get_value(row, details_col) or "").strip()
        action = str(get_value(row, action_col) or "").strip()
        cond = str(get_value(row, cond_col) or "").strip()
        category_val = get_value(row, category_col)
        category_val = (str(category_val).strip() if category_val is not None else None)
        section_val = get_value(row, section_col)
        section = (str(section_val).strip().lower() if section_val is not None else args.default_section)
        if section not in ("fault", "tcms", "emergency"):
            section = args.default_section

        cat_id = ensure_category(category_val)
        batch.append((m_id, cat_id, section, code, title, details, action, cond))

    dst.executemany(
        "INSERT INTO fault (manufacturer_id, category_id, section, code, title, details, action, condition) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
        batch,
    )

    # bump meta version
    dst.execute("UPDATE meta SET version = version + 1, updated_at = datetime('now') WHERE id = 1")
    if dst.total_changes == 0:
        dst.execute("INSERT OR REPLACE INTO meta (id, version, updated_at) VALUES (1, 1, datetime('now'))")

    dst.commit()
    src.close()
    dst.close()
    print(f"Imported {len(batch)} rows into {dest_path} for manufacturer '{args.manufacturer.lower()}'.")


if __name__ == "__main__":
    main()