# -*- coding: utf-8 -*-
"""Initialize DWS schema (dws)."""

from __future__ import annotations

from pathlib import Path
from typing import Any

from tasks.base_task import BaseTask, TaskContext


class InitDwsSchemaTask(BaseTask):
    """Apply DWS schema SQL."""

    def get_task_code(self) -> str:
        return "INIT_DWS_SCHEMA"

    def extract(self, context: TaskContext) -> dict[str, Any]:
        base_dir = Path(__file__).resolve().parents[1] / "database"
        dws_path = Path(self.config.get("schema.dws_file", base_dir / "schema_dws.sql"))
        if not dws_path.exists():
            raise FileNotFoundError(f"未找到 DWS schema 文件: {dws_path}")
        drop_first = bool(self.config.get("dws.drop_schema_first", False))
        return {"dws_sql": dws_path.read_text(encoding="utf-8"), "dws_file": str(dws_path), "drop_first": drop_first}

    def load(self, extracted: dict[str, Any], context: TaskContext) -> dict:
        with self.db.conn.cursor() as cur:
            if extracted["drop_first"]:
                cur.execute("DROP SCHEMA IF EXISTS dws CASCADE;")
                self.logger.info("已执行 DROP SCHEMA dws CASCADE")
            self.logger.info("执行 DWS schema 文件: %s", extracted["dws_file"])
            cur.execute(extracted["dws_sql"])
        return {"executed": 1, "files": [extracted["dws_file"]]}

