playwrite/db/discrete_material_plan_dao.py

"""
Data Access Object for DiscreteMaterialPlanData table.

This module provides CRUD operations for persisting discrete material plan
data to SQL Server database. It handles mapping between Chinese DataFrame
columns (from ExcelConverter) and English database columns.
"""

from db.connection import get_connection
from typing import List, Dict, Any
import pandas as pd


class DiscreteMaterialPlanDAO:
    """Data Access Object for DiscreteMaterialPlanData table"""

    def __init__(self):
        self.db = None

    def __enter__(self):
        """Enter context manager and establish database connection"""
        self.db = get_connection()
        self.db.connect()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Exit context manager and close database connection"""
        if self.db:
            self.db.disconnect()

    def close(self):
        """Close database connection"""
        if self.db:
            self.db.disconnect()

    def save_dataframe_with_replace(self, df: pd.DataFrame) -> Dict[str, int]:
        """
        Save DataFrame using REPLACE strategy (DELETE + INSERT).

        This method implements a replace strategy where existing records
        matching the plan numbers in the DataFrame are deleted before
        inserting new records.

        Args:
            df: DataFrame with discrete material plan data (Chinese column names)

        Returns:
            Dictionary with 'deleted' and 'inserted' counts

        Example:
            >>> dao = DiscreteMaterialPlanDAO()
            >>> with dao:
            ...     stats = dao.save_dataframe_with_replace(df)
            ...     print(f"Deleted: {stats['deleted']}, Inserted: {stats['inserted']}")
        """
        if df.empty:
            return {'deleted': 0, 'inserted': 0}

        # Remove duplicates based on PlanNumber and SequenceNumber
        original_count = len(df)
        df = df.drop_duplicates(subset=['备料计划单号', '序号'], keep='first')
        duplicates_removed = original_count - len(df)

        if duplicates_removed > 0:
            print(f"[INFO] 检测到 {duplicates_removed} 条重复记录（相同计划单号和序号），已自动去重")

        with get_connection() as db:
            # Get unique plan numbers
            plan_numbers = df['备料计划单号'].unique().tolist()

            # Delete existing records
            deleted = self._delete_by_plan_numbers(db, plan_numbers)

            # Insert new records in batches
            inserted = self._batch_insert(db, df)

            return {'deleted': deleted, 'inserted': inserted}

    def _delete_by_plan_numbers(self, db, plan_numbers: List[str]) -> int:
        """
        Delete records by plan numbers.

        Args:
            db: Database connection object
            plan_numbers: List of plan numbers to delete

        Returns:
            Number of records deleted
        """
        if not plan_numbers:
            return 0

        # SQL Server has a limit on IN clause parameters
        # Delete in batches to avoid exceeding the limit
        batch_size = 1000  # Safe limit for IN clause
        total_deleted = 0

        for i in range(0, len(plan_numbers), batch_size):
            batch = plan_numbers[i:i + batch_size]
            placeholders = ','.join(['?' for _ in batch])
            sql = f"DELETE FROM [dbo].[DiscreteMaterialPlanData] WHERE PlanNumber IN ({placeholders})"
            deleted = db.execute_update(sql, tuple(batch))
            total_deleted += deleted

        return total_deleted

    def _batch_insert(self, db, df: pd.DataFrame, batch_size: int = 72) -> int:
        """
        Batch insert records (max 72 per batch due to SQL Server 2100 param limit).

        SQL Server has a limit of 2100 parameters per query. With 29 fields,
        the maximum batch size is floor(2100 / 29) = 72 records per batch.

        Args:
            db: Database connection object
            df: DataFrame to insert
            batch_size: Number of records per batch (default: 72)

        Returns:
            Total number of records inserted
        """
        sql = """
        INSERT INTO [dbo].[DiscreteMaterialPlanData] (
            Factory, MaterialStatus, PlanNumber, SourceNumber, MaterialType,
            ProductCode, ProductName, ProductUnit, ProductPlanQuantity,
            UseDepartment, Remark, Creator, CreateDate, Approver, ApproveDate,
            SequenceNumber, MaterialCode, MaterialName, Specification, Model,
            DrawingNumber, MaterialQuality, PlanQuantity, Unit, RequiredDate,
            Warehouse, UnitUsage, CumulativeOutputQuantity, BOMVersion
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
        """

        total_inserted = 0
        records = self._convert_df_to_records(df)

        for i in range(0, len(records), batch_size):
            batch = records[i:i + batch_size]
            for record in batch:
                db.execute_update(sql, record)
                total_inserted += 1

        return total_inserted

    def _convert_df_to_records(self, df: pd.DataFrame) -> List[tuple]:
        """
        Convert DataFrame to list of tuples for batch insert.

        Maps Chinese DataFrame column names to English database column names
        and converts each row to a tuple in the correct order.

        Handles NaN/None values by converting them to None for NULL fields.

        Args:
            df: DataFrame with Chinese column names

        Returns:
            List of tuples, one per record
        """
        # Column order must match INSERT statement
        column_order = [
            '工厂', '备料状态', '备料计划单号', '来源单号', '备料类型', '产品编码',
            '产品名称', '产品单位', '产品计划数量', '用料部门', '备注', '制单人',
            '制单日期', '审批人', '审批日期', '序号', '材料编码', '材料名称',
            '规格', '型号', '图号', '物料材质', '计划数量', '单位', '需用日期',
            '发料仓库', '单位用量', '累计出库数量', 'BOM版本'
        ]

        # Numeric columns with their default values and data types
        numeric_columns = {
            '产品计划数量': (0, int),
            '序号': (0, int),
            '计划数量': (0, int),
            '单位用量': (0.0, float),
            '累计出库数量': (0, int),
        }

        records = []
        for _, row in df.iterrows():
            record = []
            for col in column_order:
                value = row.get(col)
                # Handle NaN, None, or empty string values
                if pd.isna(value) or value is None or (isinstance(value, str) and value.strip() == ''):
                    if col in numeric_columns:
                        # Use default value for numeric columns
                        record.append(numeric_columns[col][0])
                    else:
                        # Use None (NULL) for string columns
                        record.append(None)
                else:
                    # Convert numeric columns to proper type
                    if col in numeric_columns:
                        try:
                            default_value, target_type = numeric_columns[col]
                            if target_type == float:
                                record.append(float(value))
                            else:
                                record.append(int(value))
                        except (ValueError, TypeError):
                            # If conversion fails, use default value
                            record.append(numeric_columns[col][0])
                    else:
                        # Keep string columns as is
                        record.append(value)
            records.append(tuple(record))

        return records

    def query_by_plan_number(self, plan_number: str) -> List[Dict]:
        """
        Query all records for a specific plan number.

        Args:
            plan_number: Plan number to query

        Returns:
            List of dictionaries representing records
        """
        with get_connection() as db:
            sql = "SELECT * FROM [dbo].[DiscreteMaterialPlanData] WHERE PlanNumber = ?"
            return db.execute_query(sql, (plan_number,))

    def query_by_plan_numbers(self, plan_numbers: List[str]) -> List[Dict]:
        """
        Query records for multiple plan numbers.

        Args:
            plan_numbers: List of plan numbers to query

        Returns:
            List of dictionaries representing records
        """
        if not plan_numbers:
            return []
        placeholders = ','.join(['?' for _ in plan_numbers])
        sql = f"SELECT * FROM [dbo].[DiscreteMaterialPlanData] WHERE PlanNumber IN ({placeholders})"
        with get_connection() as db:
            return db.execute_query(sql, tuple(plan_numbers))

    def query_by_production_order(self, order_id: str) -> List[Dict]:
        """
        Query all records for a specific production order.

        Args:
            order_id: Production order ID (SourceNumber)

        Returns:
            List of dictionaries representing records
        """
        with get_connection() as db:
            sql = "SELECT * FROM [dbo].[DiscreteMaterialPlanData] WHERE SourceNumber = ?"
            return db.execute_query(sql, (order_id,))

    def count_by_plan_number(self, plan_number: str) -> int:
        """
        Count records for a specific plan number.

        Args:
            plan_number: Plan number to count

        Returns:
            Number of records
        """
        with get_connection() as db:
            sql = "SELECT COUNT(*) as count FROM [dbo].[DiscreteMaterialPlanData] WHERE PlanNumber = ?"
            result = db.execute_query(sql, (plan_number,))
            return result[0]['count'] if result else 0

    def count_all(self) -> int:
        """
        Count all records in the table.

        Returns:
            Total number of records
        """
        with get_connection() as db:
            sql = "SELECT COUNT(*) as count FROM [dbo].[DiscreteMaterialPlanData]"
            result = db.execute_query(sql)
            return result[0]['count'] if result else 0

    def delete_by_plan_numbers(self, plan_numbers: List[str]) -> int:
        """
        Delete all records for specified plan numbers.

        Args:
            plan_numbers: List of plan numbers to delete

        Returns:
            Number of records deleted
        """
        with get_connection() as db:
            return self._delete_by_plan_numbers(db, plan_numbers)

    def get_statistics(self) -> Dict[str, Any]:
        """
        Get comprehensive statistics about the data.

        Returns:
            Dictionary with statistics including total records,
            unique plans, unique orders, and date range
        """
        with get_connection() as db:
            sql = """
            SELECT
                COUNT(*) as total_records,
                COUNT(DISTINCT PlanNumber) as unique_plans,
                COUNT(DISTINCT SourceNumber) as unique_orders,
                MIN(CreateDate) as earliest_record,
                MAX(CreateDate) as latest_record
            FROM [dbo].[DiscreteMaterialPlanData]
            """
            result = db.execute_query(sql)
            return result[0] if result else {}

    # ==================== ENHANCED QUERY METHODS ====================

    def query_all(self) -> List[Dict]:
        """
        Query all records from DiscreteMaterialPlanData table.

        Returns:
            List of dictionaries representing all records
        """
        with get_connection() as db:
            sql = "SELECT * FROM [dbo].[DiscreteMaterialPlanData]"
            return db.execute_query(sql)

    def query_by_source_numbers(self, source_numbers: List[str]) -> List[Dict]:
        """
        Query records by SourceNumber list (生产订单号).

        Args:
            source_numbers: List of SourceNumber values to query

        Returns:
            List of dictionaries representing records
        """
        if not source_numbers:
            return []

        # SQL Server parameter limit requires batching
        batch_size = 2000
        all_results = []

        for i in range(0, len(source_numbers), batch_size):
            batch = source_numbers[i:i + batch_size]
            placeholders = ','.join(['?' for _ in batch])
            sql = f"SELECT * FROM [dbo].[DiscreteMaterialPlanData] WHERE SourceNumber IN ({placeholders})"
            with get_connection() as db:
                results = db.execute_query(sql, tuple(batch))
                all_results.extend(results)

        return all_results

    def get_unique_material_names(self, source_numbers: List[str] = None) -> List[str]:
        """
        Get unique material names, optionally filtered by SourceNumber.

        Args:
            source_numbers: Optional list of SourceNumber values to filter by

        Returns:
            List of unique material names
        """
        if source_numbers is None or not source_numbers:
            # No filter - get all unique material names
            sql = "SELECT DISTINCT MaterialName FROM [dbo].[DiscreteMaterialPlanData] WHERE MaterialName IS NOT NULL"
            with get_connection() as db:
                results = db.execute_query(sql)
                return [r['MaterialName'] for r in results if r.get('MaterialName')]
        else:
            # Filter by SourceNumber list
            batch_size = 2000
            all_material_names = set()

            for i in range(0, len(source_numbers), batch_size):
                batch = source_numbers[i:i + batch_size]
                placeholders = ','.join(['?' for _ in batch])
                sql = f"""
                SELECT DISTINCT MaterialName
                FROM [dbo].[DiscreteMaterialPlanData]
                WHERE SourceNumber IN ({placeholders})
                  AND MaterialName IS NOT NULL
                """
                with get_connection() as db:
                    results = db.execute_query(sql, tuple(batch))
                    batch_materials = [r['MaterialName'] for r in results if r.get('MaterialName')]
                    all_material_names.update(batch_materials)

            return list(all_material_names)

    def get_material_names_by_总排号(self, 总排号_list: List[str]) -> List[str]:
        """
        Get unique material names by 总排号 list.
        This method combines query from production contract data and discrete material plan.

        Args:
            总排号_list: List of 总排号 values

        Returns:
            List of unique material names
        """
        from db.production_contract_data_dao import ProductionContractDataDAO

        # First get SourceNumbers from production contract data
        contract_dao = ProductionContractDataDAO()
        source_numbers = contract_dao.get_source_numbers_by_总排号(总排号_list)

        # Then get material names filtered by these SourceNumbers
        return self.get_unique_material_names(source_numbers)