from __future__ import annotations import json import logging from pathlib import Path from typing import Any logger = logging.getLogger(__name__) class ErpTakeoffHeaderMixin: def create_drawing_takeoff(self, payload: dict[str, Any]) -> dict[str, Any]: pages_data = payload.get("pages_data") if isinstance(pages_data, (list, dict)): pages_data_json = json.dumps(pages_data) else: pages_data_json = str(pages_data) if pages_data else "[]" defaults = { "id": payload.get("id"), "project_id": payload.get("project_id"), "document_id": payload.get("document_id"), "drawing_code": payload.get("drawing_code", "DWG-001"), "drawing_title": payload.get( "drawing_title", payload.get("drawing_code", "Bản vẽ kỹ thuật") ), "drawing_type": payload.get("drawing_type", "civil_building"), "drawing_scale": payload.get("drawing_scale", "1:100"), "file_url": payload.get("file_url", ""), "file_path": payload.get("file_path", ""), "file_hash": payload.get("file_hash", ""), "preview_image_url": payload.get("preview_image_url", ""), "total_estimated_cost_vnd": payload.get("total_estimated_cost_vnd", 0.0), "total_concrete_volume_m3": payload.get("total_concrete_volume_m3", 0.0), "total_formwork_area_m2": payload.get("total_formwork_area_m2", 0.0), "total_rebar_weight_tons": payload.get("total_rebar_weight_tons", 0.0), "total_earthwork_volume_m3": payload.get("total_earthwork_volume_m3", 0.0), "takeoff_status": payload.get("takeoff_status", "pending"), "ai_vision_model": payload.get("ai_vision_model", "gemini-3.1-pro-high"), "ai_analysis_notes": payload.get("ai_analysis_notes", ""), "confidence_score": int(payload.get("confidence_score", 98)), "confidence_level": payload.get("confidence_level", "hsmt_verified"), "hsmt_reference_file": payload.get("hsmt_reference_file", None), "ai_report_message": payload.get( "ai_report_message", "AI Quỳnh QS đã sẵn sàng đối soát bóc tách khối lượng.", ), "pages_data": pages_data_json, } with self.get_connection() as conn: with conn.cursor() as cur: sql = """ INSERT INTO erp_drawing_takeoffs ( id, project_id, document_id, drawing_code, drawing_title, drawing_type, drawing_scale, file_url, file_path, file_hash, preview_image_url, total_estimated_cost_vnd, total_concrete_volume_m3, total_formwork_area_m2, total_rebar_weight_tons, total_earthwork_volume_m3, takeoff_status, ai_vision_model, ai_analysis_notes, confidence_score, confidence_level, hsmt_reference_file, ai_report_message, pages_data ) VALUES ( %(id)s, %(project_id)s, %(document_id)s, %(drawing_code)s, %(drawing_title)s, %(drawing_type)s, %(drawing_scale)s, %(file_url)s, %(file_path)s, %(file_hash)s, %(preview_image_url)s, %(total_estimated_cost_vnd)s, %(total_concrete_volume_m3)s, %(total_formwork_area_m2)s, %(total_rebar_weight_tons)s, %(total_earthwork_volume_m3)s, %(takeoff_status)s, %(ai_vision_model)s, %(ai_analysis_notes)s, %(confidence_score)s, %(confidence_level)s, %(hsmt_reference_file)s, %(ai_report_message)s, %(pages_data)s ) RETURNING *; """ cur.execute(sql, defaults) conn.commit() return cur.fetchone() def get_drawing_takeoff(self, takeoff_id: str) -> dict[str, Any] | None: """Lấy thông tin chi tiết một hồ sơ bóc tách bản vẽ.""" sql = """ SELECT t.*, p.project_name, p.project_code FROM erp_drawing_takeoffs t LEFT JOIN projects p ON t.project_id = p.id WHERE t.id = %s; """ with self.get_connection() as conn, conn.cursor() as cur: cur.execute(sql, (takeoff_id,)) return cur.fetchone() def list_drawing_takeoffs( self, project_id: str | None = None, limit: int = 50 ) -> list[dict[str, Any]]: """Lấy danh sách các hồ sơ bóc tách bản vẽ.""" if project_id: sql = """ SELECT t.*, p.project_name, p.project_code FROM erp_drawing_takeoffs t LEFT JOIN projects p ON t.project_id = p.id WHERE t.project_id = %s ORDER BY t.created_at DESC LIMIT %s; """ params = (project_id, limit) else: sql = """ SELECT t.*, p.project_name, p.project_code FROM erp_drawing_takeoffs t LEFT JOIN projects p ON t.project_id = p.id ORDER BY t.created_at DESC LIMIT %s; """ params = (limit,) with self.get_connection() as conn, conn.cursor() as cur: cur.execute(sql, params) return cur.fetchall() def update_drawing_takeoff( self, takeoff_id: str, payload: dict[str, Any] ) -> dict[str, Any] | None: """Cập nhật thông tin hồ sơ bóc tách bản vẽ (trạng thái, tổng khối lượng, dự toán, ghi chú).""" if not payload: return None set_clauses = [] params = [] for k, v in payload.items(): set_clauses.append(f"{k} = %s") if k == "pages_data" and isinstance(v, (list, dict)): params.append(json.dumps(v)) else: params.append(v) params.append(takeoff_id) sql = f"UPDATE erp_drawing_takeoffs SET {', '.join(set_clauses)}, updated_at = CURRENT_TIMESTAMP WHERE id = %s RETURNING *;" with self.get_connection() as conn, conn.cursor() as cur: cur.execute(sql, params) row = cur.fetchone() conn.commit() return row def delete_drawing_takeoff(self, takeoff_id: str) -> bool: """Xóa hồ sơ bóc tách bản vẽ và toàn bộ các dòng chi tiết liên quan.""" with self.get_connection() as conn: with conn.cursor() as cur: # Xóa items trước để đảm bảo sạch sẽ cur.execute( "DELETE FROM erp_drawing_takeoff_items WHERE takeoff_id = %s;", (takeoff_id,), ) # Xóa bản ghi header cur.execute( "DELETE FROM erp_drawing_takeoffs WHERE id = %s RETURNING id, file_path;", (takeoff_id,), ) deleted = cur.fetchone() conn.commit() if deleted and deleted.get("file_path"): try: p = Path(deleted["file_path"]) is_safe_temp_dir = any( part in p.parts for part in ["uploads", "tmp", "temp", "secure_vault"] ) is_repo_source = any( part in p.parts for part in ["CongTrinh", "HĐ-2026", "HĐ-2025", "HĐ-2024", "HĐ-2023", "HĐ-2022", "KNOWLEDGE_BASE"] ) if is_safe_temp_dir and not is_repo_source and p.exists(): p.unlink(missing_ok=True) except Exception: pass # Clean up generated images in static/drawings drawings_dir = Path(__file__).resolve().parents[4] / "static" / "drawings" if drawings_dir.exists(): for f in drawings_dir.glob(f"{takeoff_id}*.*"): try: f.unlink() except Exception: pass for f in drawings_dir.glob(f"{p.stem}*.*"): try: f.unlink() except Exception: pass except Exception as e: logger.warning( "[TAKEOFF] Không thể xóa file vật lý %s: %s", deleted["file_path"], e, ) return bool(deleted) def delete_all_drawing_takeoffs(self, project_id: str | None = None) -> int: """Xóa hàng loạt hồ sơ bóc tách bản vẽ cũ (Dành riêng cho Super Admin).""" with self.get_connection() as conn: with conn.cursor() as cur: if project_id: cur.execute( """ DELETE FROM erp_drawing_takeoff_items WHERE takeoff_id IN (SELECT id FROM erp_drawing_takeoffs WHERE project_id = %s); """, (project_id,), ) cur.execute( "DELETE FROM erp_drawing_takeoffs WHERE project_id = %s RETURNING id;", (project_id,), ) else: cur.execute("DELETE FROM erp_drawing_takeoff_items;") cur.execute("DELETE FROM erp_drawing_takeoffs RETURNING id;") rows = cur.fetchall() conn.commit() return len(rows)