"""物流公司数据访问类"""
import pandas as pd
from typing import Dict, List, Optional
from dataaccess.base_dao import BaseDAO


class CompanyDAO(BaseDAO):
    """物流公司数据访问"""

    def __init__(self):
        super().__init__("logistics_company")

    def get_all_companies(self, country: str = None, active_only: bool = True) -> List[Dict]:
        """获取所有物流公司"""
        conditions = {}
        if country:
            conditions["country"] = country
        if active_only:
            conditions["active"] = 1

        if conditions:
            df = self.get_by_condition(conditions)
        else:
            df = self.get_all()

        if df.empty:
            return []

        return df.to_dict("records")

    def get_company_info(self, company_code: str) -> Optional[Dict]:
        """获取物流公司详细信息"""
        df = self.get_by_condition({"company_code": company_code})
        if df.empty:
            return None
        return df.iloc[0].to_dict()

    def get_companies_by_type(self, country: str, logistics_type: str) -> List[Dict]:
        """根据物流类型获取公司列表"""
        query = f"""
            SELECT * FROM {self.table_name}
            WHERE country = '{country}'
            AND logistics_type = '{logistics_type}'
            AND active = 1
        """
        df = self.execute_query(query)
        if df.empty:
            return []
        return df.to_dict("records")