"""
护理项目收入表数据处理脚本

功能：从源表 dws_adv_xmsrb_hl 读取数据，按追加模式写入目标表
源表：dws_adv_xmsrb_hl（项目收入表-护理）
目标表：dws_adv_xmsrb_hl（同名表，追加模式）
更新模式：append（追加）

作者：cursor
创建时间：2025-12-31
"""

import argparse
import logging
import os
import sys
from datetime import datetime
from typing import Any, Dict, List, Optional

from sqlalchemy import create_engine, inspect, text
from sqlalchemy.engine import Engine
from sqlalchemy.orm import Session, sessionmaker

# 添加项目根目录到路径
sys.path.insert(
    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
)

try:
    from app.config.config import config, current_env  # type: ignore

    # 根据当前环境获取配置类
    Config = config.get(current_env, config["default"])
except ImportError:
    # 如果无法导入，使用环境变量
    class Config:  # type: ignore
        SQLALCHEMY_DATABASE_URI = os.environ.get(
            "DATABASE_URI", "postgresql://user:password@localhost:5432/database"
        )


# 配置日志
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)


class NursingProjectIncomeProcessor:
    """护理项目收入表数据处理器"""

    # 源表名称
    SOURCE_TABLE = "dws_adv_xmsrb_hl"

    # 目标表名称（与源表相同，追加模式）
    TARGET_TABLE = "dws_adv_xmsrb_hl"

    # 源表字段定义
    SOURCE_COLUMNS = [
        "srje",  # 收入金额 numeric
        "sfcs",  # 收费次数 numeric
        "dim_key",  # 维度key varchar
        "sr",  # 收入 numeric
        "ksdm",  # 科室代码 varchar
        "sfbm",  # 收费编码 varchar
        "ny",  # 年月 timestamp
    ]

    def __init__(
        self,
        source_schema: str = "public",
        target_schema: str = "public",
        update_mode: str = "append",
        batch_size: int = 1000,
    ):
        """
        初始化处理器

        Args:
            source_schema: 源表schema
            target_schema: 目标表schema
            update_mode: 更新模式，'append'（追加）或 'full'（全量更新）
            batch_size: 批量处理大小
        """
        self.source_schema = source_schema
        self.target_schema = target_schema
        self.update_mode = update_mode.lower()
        self.batch_size = batch_size

        self.engine: Optional[Engine] = None
        self.session: Optional[Session] = None

        self.processed_count = 0
        self.inserted_count = 0
        self.error_count = 0

        # 验证更新模式
        if self.update_mode not in ["append", "full"]:
            raise ValueError(
                f"不支持的更新模式: {update_mode}，仅支持 'append' 或 'full'"
            )

        logger.info(
            f"初始化护理项目收入表处理器: "
            f"源表={source_schema}.{self.SOURCE_TABLE}, "
            f"目标表={target_schema}.{self.TARGET_TABLE}, "
            f"更新模式={update_mode}"
        )

    def connect_database(self) -> bool:
        """
        连接数据库

        Returns:
            连接是否成功
        """
        try:
            db_uri = Config.SQLALCHEMY_DATABASE_URI

            if not db_uri:
                logger.error("未找到数据库配置（SQLALCHEMY_DATABASE_URI）")
                return False

            self.engine = create_engine(db_uri)
            SessionLocal = sessionmaker(bind=self.engine)
            self.session = SessionLocal()

            # 测试连接
            with self.engine.connect() as conn:
                conn.execute(text("SELECT 1"))

            logger.info(f"成功连接数据库: {db_uri.split('@')[-1]}")
            return True

        except Exception as e:
            logger.error(f"连接数据库失败: {str(e)}")
            return False

    def check_table_exists(self, table_name: str, schema: str = "public") -> bool:
        """
        检查表是否存在

        Args:
            table_name: 表名
            schema: schema名

        Returns:
            表是否存在
        """
        try:
            if not self.engine:
                return False

            inspector = inspect(self.engine)
            tables = inspector.get_table_names(schema=schema)
            return table_name in tables

        except Exception as e:
            logger.error(f"检查表是否存在失败: {str(e)}")
            return False

    def create_target_table_if_not_exists(self) -> bool:
        """
        如果目标表不存在，则创建

        Returns:
            操作是否成功
        """
        try:
            if self.check_table_exists(self.TARGET_TABLE, self.target_schema):
                logger.info(f"目标表 {self.target_schema}.{self.TARGET_TABLE} 已存在")
                return True

            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            # 创建目标表 DDL
            create_sql = text(f"""
                CREATE TABLE IF NOT EXISTS {self.target_schema}.{self.TARGET_TABLE} (
                    id SERIAL PRIMARY KEY,
                    srje NUMERIC COMMENT '收入金额',
                    sfcs NUMERIC COMMENT '收费次数',
                    dim_key VARCHAR(255) COMMENT '维度key',
                    sr NUMERIC COMMENT '收入',
                    ksdm VARCHAR(100) COMMENT '科室代码',
                    sfbm VARCHAR(100) COMMENT '收费编码',
                    ny TIMESTAMP COMMENT '年月',
                    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP COMMENT '数据创建时间'
                );
                COMMENT ON TABLE {self.target_schema}.{self.TARGET_TABLE} IS '项目收入表-护理';
            """)

            self.session.execute(create_sql)
            self.session.commit()

            logger.info(f"成功创建目标表 {self.target_schema}.{self.TARGET_TABLE}")
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"创建目标表失败: {str(e)}")
            return False

    def extract_source_data(
        self,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> List[Dict[str, Any]]:
        """
        从源表提取数据

        Args:
            start_date: 开始日期（过滤 ny 字段）
            end_date: 结束日期（过滤 ny 字段）
            limit: 限制提取的数据行数

        Returns:
            数据行列表
        """
        try:
            if not self.session:
                logger.error("数据库会话未初始化")
                return []

            # 构建查询
            columns_str = ", ".join(self.SOURCE_COLUMNS)
            query = (
                f"SELECT {columns_str} FROM {self.source_schema}.{self.SOURCE_TABLE}"
            )

            # 添加日期过滤条件
            conditions = []
            params: Dict[str, Any] = {}

            if start_date:
                conditions.append("ny >= :start_date")
                params["start_date"] = start_date

            if end_date:
                conditions.append("ny <= :end_date")
                params["end_date"] = end_date

            if conditions:
                query += " WHERE " + " AND ".join(conditions)

            # 添加排序
            query += " ORDER BY ny DESC"

            # 添加限制
            if limit:
                query += f" LIMIT {limit}"

            logger.info(f"执行查询: {query}")
            result = self.session.execute(text(query), params)

            # 转换为字典列表
            rows = []
            for row in result:
                row_dict = dict(zip(self.SOURCE_COLUMNS, row))
                rows.append(row_dict)

            self.processed_count = len(rows)
            logger.info(f"从源表提取了 {self.processed_count} 条数据")
            return rows

        except Exception as e:
            logger.error(f"提取源数据失败: {str(e)}")
            return []

    def clear_target_table(self) -> bool:
        """
        清空目标表（用于全量更新模式）

        Returns:
            操作是否成功
        """
        try:
            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            delete_sql = text(f"DELETE FROM {self.target_schema}.{self.TARGET_TABLE}")
            self.session.execute(delete_sql)
            self.session.commit()

            logger.info(f"目标表 {self.target_schema}.{self.TARGET_TABLE} 已清空")
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"清空目标表失败: {str(e)}")
            return False

    def insert_data(self, data_rows: List[Dict[str, Any]]) -> bool:
        """
        将数据插入目标表

        Args:
            data_rows: 数据行列表

        Returns:
            插入是否成功
        """
        try:
            if not data_rows:
                logger.warning("没有数据需要插入")
                return True

            if not self.session:
                logger.error("数据库会话未初始化")
                return False

            # 全量更新模式：先清空目标表
            if self.update_mode == "full" and not self.clear_target_table():
                return False

            # 构建插入 SQL
            columns_str = ", ".join(self.SOURCE_COLUMNS + ["create_time"])
            placeholders = ", ".join(
                [f":{col}" for col in self.SOURCE_COLUMNS] + ["CURRENT_TIMESTAMP"]
            )

            insert_sql = text(f"""
                INSERT INTO {self.target_schema}.{self.TARGET_TABLE} ({columns_str})
                VALUES ({placeholders})
            """)

            # 批量插入
            success_count = 0
            for i, row in enumerate(data_rows):
                try:
                    self.session.execute(insert_sql, row)
                    success_count += 1

                    # 批量提交
                    if success_count % self.batch_size == 0:
                        self.session.commit()
                        logger.info(f"已插入 {success_count} 条数据...")

                except Exception as e:
                    self.error_count += 1
                    logger.error(f"插入数据失败 (行 {i}): {str(e)}")

            # 最终提交
            self.session.commit()
            self.inserted_count = success_count

            logger.info(
                f"数据插入完成: 成功 {self.inserted_count} 条, 失败 {self.error_count} 条"
            )
            return True

        except Exception as e:
            if self.session:
                self.session.rollback()
            logger.error(f"批量插入数据失败: {str(e)}")
            return False

    def close_connection(self):
        """关闭数据库连接"""
        if self.session:
            try:
                self.session.close()
                logger.info("数据库会话已关闭")
            except Exception as e:
                logger.error(f"关闭数据库会话失败: {str(e)}")

        if self.engine:
            try:
                self.engine.dispose()
                logger.info("数据库引擎已释放")
            except Exception as e:
                logger.error(f"释放数据库引擎失败: {str(e)}")

    def run(
        self,
        start_date: Optional[str] = None,
        end_date: Optional[str] = None,
        limit: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        执行数据处理流程

        Args:
            start_date: 开始日期（过滤 ny 字段）
            end_date: 结束日期（过滤 ny 字段）
            limit: 限制处理的数据行数

        Returns:
            执行结果
        """
        result = {
            "success": False,
            "processed_count": 0,
            "inserted_count": 0,
            "error_count": 0,
            "update_mode": self.update_mode,
            "message": "",
            "start_time": datetime.now().isoformat(),
            "end_time": None,
        }

        try:
            logger.info("=" * 60)
            logger.info("开始护理项目收入表数据处理")
            logger.info(f"源表: {self.source_schema}.{self.SOURCE_TABLE}")
            logger.info(f"目标表: {self.target_schema}.{self.TARGET_TABLE}")
            logger.info(f"更新模式: {self.update_mode}")
            if start_date:
                logger.info(f"开始日期: {start_date}")
            if end_date:
                logger.info(f"结束日期: {end_date}")
            logger.info("=" * 60)

            # 1. 连接数据库
            if not self.connect_database():
                result["message"] = "连接数据库失败"
                return result

            # 2. 检查/创建目标表
            if not self.create_target_table_if_not_exists():
                result["message"] = "创建目标表失败"
                return result

            # 3. 提取源数据
            data_rows = self.extract_source_data(
                start_date=start_date,
                end_date=end_date,
                limit=limit,
            )

            if not data_rows:
                result["message"] = "未提取到数据"
                result["success"] = True  # 没有数据不算失败
                return result

            # 4. 插入数据到目标表
            if self.insert_data(data_rows):
                result["success"] = True
                result["processed_count"] = self.processed_count
                result["inserted_count"] = self.inserted_count
                result["error_count"] = self.error_count
                result["message"] = (
                    f"处理完成: 成功 {self.inserted_count} 条, "
                    f"失败 {self.error_count} 条"
                )
            else:
                result["message"] = "插入数据失败"

        except Exception as e:
            logger.error(f"处理过程发生异常: {str(e)}")
            result["message"] = f"处理失败: {str(e)}"

        finally:
            result["end_time"] = datetime.now().isoformat()
            self.close_connection()

        logger.info("=" * 60)
        logger.info(f"处理结果: {result['message']}")
        logger.info("=" * 60)

        return result


def process_nursing_project_income(
    source_schema: str = "public",
    target_schema: str = "public",
    update_mode: str = "append",
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    limit: Optional[int] = None,
) -> Dict[str, Any]:
    """
    处理护理项目收入表数据（入口函数）

    Args:
        source_schema: 源表schema
        target_schema: 目标表schema
        update_mode: 更新模式，'append'（追加）或 'full'（全量更新）
        start_date: 开始日期（过滤 ny 字段）
        end_date: 结束日期（过滤 ny 字段）
        limit: 限制处理的数据行数

    Returns:
        处理结果
    """
    processor = NursingProjectIncomeProcessor(
        source_schema=source_schema,
        target_schema=target_schema,
        update_mode=update_mode,
    )
    return processor.run(
        start_date=start_date,
        end_date=end_date,
        limit=limit,
    )


def parse_args():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="护理项目收入表数据处理工具")

    parser.add_argument(
        "--source-schema",
        type=str,
        default="public",
        help="源表schema（默认：public）",
    )

    parser.add_argument(
        "--target-schema",
        type=str,
        default="public",
        help="目标表schema（默认：public）",
    )

    parser.add_argument(
        "--update-mode",
        type=str,
        choices=["append", "full"],
        default="append",
        help="更新模式：append（追加）或 full（全量更新），默认：append",
    )

    parser.add_argument(
        "--start-date",
        type=str,
        default=None,
        help="开始日期，格式：YYYY-MM-DD（过滤 ny 字段）",
    )

    parser.add_argument(
        "--end-date",
        type=str,
        default=None,
        help="结束日期，格式：YYYY-MM-DD（过滤 ny 字段）",
    )

    parser.add_argument(
        "--limit",
        type=int,
        default=None,
        help="限制处理的数据行数",
    )

    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()

    result = process_nursing_project_income(
        source_schema=args.source_schema,
        target_schema=args.target_schema,
        update_mode=args.update_mode,
        start_date=args.start_date,
        end_date=args.end_date,
        limit=args.limit,
    )

    # 输出结果
    print("\n" + "=" * 60)
    print(f"处理结果: {'成功' if result['success'] else '失败'}")
    print(f"消息: {result['message']}")
    print(f"处理: {result['processed_count']} 条")
    print(f"插入: {result['inserted_count']} 条")
    print(f"失败: {result['error_count']} 条")
    print(f"更新模式: {result['update_mode']}")
    print(f"开始时间: {result['start_time']}")
    print(f"结束时间: {result['end_time']}")
    print("=" * 60)

    # 设置退出代码
    exit(0 if result["success"] else 1)
