"""
Data Access Object for production contract data.

This module provides query operations for accessing production contract data
from the [productionContractData].[26年压力表合同数据] table.
"""

from typing import List, Dict, Any
from db.connection import get_connection


class ProductionContractDataDAO:
    """Data Access Object for production contract data queries"""

    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 query_by_总排号(self, 总排号_list: List[str]) -> List[Dict[str, Any]]:
        """
        Query production contract data by 总排号 list.

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

        Returns:
            List of dictionaries containing 总排号, 生产订单号, 序号, 订单号, 客户名称, 产品型号
        """
        if not 总排号_list:
            return []

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

        for i in range(0, len(总排号_list), batch_size):
            batch = 总排号_list[i:i + batch_size]
            placeholders = ','.join(['?' for _ in batch])
            sql = f"""
            SELECT [总排号], [生产订单号], [序号], [订单号], [客户名称], [产品型号]
            FROM [productionContractData].[26年压力表合同数据]
            WHERE [总排号] IN ({placeholders})
            ORDER BY [序号]
            """

            with get_connection() as db:
                results = db.execute_query(sql, tuple(batch))
                all_results.extend(results)

        return all_results

    def get_source_numbers_by_总排号(self, 总排号_list: List[str]) -> List[str]:
        """
        Extract unique 生产订单号 values by 总排号 list.

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

        Returns:
            List of unique 生产订单号 values (SourceNumber)
        """
        results = self.query_by_总排号(总排号_list)
        # Extract unique 生产订单号 values, excluding None/null values
        source_numbers = list(set(
            [r['生产订单号'] for r in results if r.get('生产订单号')]
        ))
        return source_numbers

    def get_生产订单号_map(self, 总排号_list: List[str]) -> Dict[str, str]:
        """
        Get mapping between 总排号 and 生产订单号.

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

        Returns:
            Dictionary mapping 总排号 -> 生产订单号
        """
        results = self.query_by_总排号(总排号_list)
        return {
            r['总排号']: r['生产订单号']
            for r in results
            if r.get('总排号') and r.get('生产订单号')
        }