#!/usr/bin/env python
# coding=utf-8
"""DataWorks 客户端：根据表名获取生产代码"""

import os
import time
from pathlib import Path

from odps_skills.client import ODPSClient, ODPS_CONFIGS
from odps_skills.config import get_data_dir

# DataWorks SDK（可选依赖，仅 DataWorksClient 用到）
try:
    from alibabacloud_dataworks_public20240518.client import Client as _DWClient
    from alibabacloud_tea_openapi import models as _open_api_models
    from alibabacloud_dataworks_public20240518 import models as _dw_models
    _DW_AVAILABLE = True
except ImportError:
    _DW_AVAILABLE = False

# 账号下所有可访问的 DataWorks 项目（project_id → name）
_DW_PROJECTS = {
    4858:   "loghubods",
    11300:  "DWH",
    5477:   "videocdm",
    548768: "piaoquan_api",
    148813: "content_safety",
    96094:  "algo",
    52578:  "majin",
    5057:   "useractionbi",
    5034:   "user_video_action_cdm",
    4868:   "usercdm",
    4859:   "videoods",
    6025:   "videoads",
    5535:   "user_video_tag",
    19288:  "RecallEmbedding",
    10762:  "Test_model1",
    193831: "cost_mgt_1894469520484605",
    156474: "dyp_1",
    156475: "dyp_2",
    343868: "pq_data_space",
    343957: "pq_grafana_se",
}

def _get_cache_dir():
    """缓存目录：基于可配置的 data_dir。"""
    return os.path.join(get_data_dir(), "production_code")


def _call_with_retry(fn, max_retries=3, base_delay=2):
    """带限流重试的 API 调用包装。"""
    for attempt in range(max_retries):
        try:
            return fn()
        except Exception as e:
            if "Throttling" in str(e) and attempt < max_retries - 1:
                delay = base_delay * (2 ** attempt)
                print(f"  [throttled] 等待 {delay}s 后重试...")
                time.sleep(delay)
                continue
            raise


class DataWorksClient:
    def __init__(self):
        if not _DW_AVAILABLE:
            raise ImportError(
                "请先安装 DataWorks SDK：\n"
                "pip install alibabacloud-dataworks-public20240518"
            )
        # 初始化所有 AK 对应的客户端（不同 AK 对不同项目有权限）
        self._clients = {}
        for config_name, cfg in ODPS_CONFIGS.items():
            dw_config = _open_api_models.Config(
                access_key_id=cfg["access_id"],
                access_key_secret=cfg["access_secret"],
                endpoint="dataworks.cn-hangzhou.aliyuncs.com",
            )
            self._clients[config_name] = _DWClient(dw_config)
        self.client = self._clients["default"]

    @staticmethod
    def _build_entity_id(table_name: str) -> str:
        """构造 GetTable 的 entity ID。

        支持格式：
          - project.table  → maxcompute-table:::project::table
          - table           → maxcompute-table:::loghubods::table
        """
        parts = table_name.split(".", 1)
        if len(parts) == 2:
            project, table = parts
        else:
            project, table = "loghubods", parts[0]
        return f"maxcompute-table:::{project}::{table}"

    def get_table_info(self, table_name: str) -> dict:
        """获取表的元信息（含上游任务列表）。

        Returns:
            dict with keys: name, comment, dataworks_tasks[{id, name}], ...
        """
        entity_id = self._build_entity_id(table_name)
        resp = _call_with_retry(lambda: self.client.get_table(
            _dw_models.GetTableRequest(id=entity_id, include_business_metadata=True)
        ))
        table = resp.body.to_map().get("Table", {})
        biz = table.get("BusinessMetadata", {})
        return {
            "id": table.get("Id"),
            "name": table.get("Name"),
            "comment": table.get("Comment"),
            "project_id": biz.get("Extension", {}).get("ProjectId"),
            "dataworks_tasks": biz.get("UpstreamTasks", []),
            "partition_keys": table.get("PartitionKeys", []),
        }

    def _get_task_code(self, task_id: int) -> dict:
        """尝试用所有 AK 获取任务代码，返回第一个成功的结果。"""
        for config_name, client in self._clients.items():
            try:
                resp = _call_with_retry(lambda c=client: c.get_task(
                    _dw_models.GetTaskRequest(id=task_id, project_env="Prod")
                ))
                task = resp.body.to_map().get("Task", {})
                return {
                    "task_id": task_id,
                    "task_name": task.get("Name"),
                    "task_type": task.get("Type"),
                    "content": (task.get("Script") or {}).get("Content", ""),
                    "config": config_name,
                }
            except Exception as e:
                if "11020205003" in str(e):
                    continue  # 无权限，尝试下一个 AK
                raise
        return None

    @staticmethod
    def _normalize_table_name(table_name: str) -> str:
        """补全 project 前缀：table → loghubods.table"""
        if "." not in table_name:
            return f"loghubods.{table_name}"
        return table_name

    @staticmethod
    def _cache_path(table_name: str) -> str:
        return os.path.join(_get_cache_dir(), f"{table_name}.sql")

    @staticmethod
    def _schema_cache_path(table_name: str) -> str:
        return os.path.join(_get_cache_dir(), f"{table_name}.json")

    def _read_cache(self, table_name: str) -> str | None:
        path = self._cache_path(table_name)
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return f.read()
        return None

    def _write_cache(self, table_name: str, content: str):
        os.makedirs(_get_cache_dir(), exist_ok=True)
        with open(self._cache_path(table_name), "w", encoding="utf-8") as f:
            f.write(content)

    def _read_schema_cache(self, table_name: str) -> dict | None:
        import json
        path = self._schema_cache_path(table_name)
        if os.path.exists(path):
            with open(path, "r", encoding="utf-8") as f:
                return json.load(f)
        return None

    def _write_schema_cache(self, table_name: str, schema: dict):
        import json
        os.makedirs(_get_cache_dir(), exist_ok=True)
        with open(self._schema_cache_path(table_name), "w", encoding="utf-8") as f:
            json.dump(schema, f, ensure_ascii=False, indent=2)

    def _ensure_schema_cache(self, table_name: str, force: bool = False,
                              dataworks_tasks: list | None = None):
        """确保 schema 缓存存在，无则拉取并写入。

        Args:
            dataworks_tasks: 预获取的上游任务列表，避免重复调用 get_table_info()
        """
        if not force:
            cached = self._read_schema_cache(table_name)
            if cached is not None:
                return
        try:
            schema = self.get_table_schema(table_name, dataworks_tasks=dataworks_tasks)
            self._write_schema_cache(table_name, schema)
            print(f"[saved] {self._schema_cache_path(table_name)}")
        except Exception as e:
            print(f"[WARN] 获取表结构失败 {table_name}: {e}")

    def get_table_schema(self, table_name: str,
                          dataworks_tasks: list | None = None) -> dict:
        """通过 ODPS SDK 获取表结构元信息。

        Args:
            table_name: 表名（支持 project.table 格式）
            dataworks_tasks: 预获取的上游任务列表，避免重复 API 调用

        Returns:
            dict: {name, project, comment, columns, partition_keys, dataworks_tasks}
        """
        table_name = self._normalize_table_name(table_name)
        parts = table_name.split(".", 1)
        project, table = parts[0], parts[1]

        # 用默认 AK 对应的 ODPSClient 获取 ODPS 表结构
        odps_client = ODPSClient(project=project)
        t = odps_client.odps.get_table(table)

        columns = [
            {"name": c.name, "type": str(c.type), "comment": c.comment or ""}
            for c in t.table_schema.columns
        ]
        partition_keys = [
            {"name": c.name, "type": str(c.type), "comment": c.comment or ""}
            for c in t.table_schema.partitions
        ]

        # 上游任务：优先用传入的，否则从 DataWorks API 获取
        if dataworks_tasks is None:
            try:
                info = self.get_table_info(table_name)
                dataworks_tasks = [
                    {"id": task.get("Id"), "name": task.get("Name")}
                    for task in info.get("dataworks_tasks", [])
                ]
            except Exception:
                dataworks_tasks = []

        # 直接上游表（血缘）
        try:
            upstream_tables = self.get_upstream_tables(table_name)
        except Exception:
            upstream_tables = []

        return {
            "name": table,
            "project": project,
            "comment": t.comment or "",
            "columns": columns,
            "partition_keys": partition_keys,
            "dataworks_tasks": dataworks_tasks,
            "upstream_tables": upstream_tables,
        }

    def get_node_code(self, table_name: str, force: bool = False) -> list:
        """根据表名获取生产代码（优先读本地缓存）。

        流程：本地缓存 → GetTable → GetTask → 写缓存 → 返回代码

        Args:
            table_name: 表名（支持 project.table 格式）
            force: True 时跳过缓存，强制从 API 拉取

        Returns:
            list of dict，每条包含：
                task_id, task_name, task_type, content
        """
        table_name = self._normalize_table_name(table_name)

        # 读缓存
        if not force:
            cached = self._read_cache(table_name)
            if cached is not None:
                print(f"[cache] {self._cache_path(table_name)}")
                # 同时检查 schema 缓存，无则补拉
                self._ensure_schema_cache(table_name, force=False)
                return [{"task_id": None, "task_name": "(cached)", "task_type": None, "content": cached}]

        # API 拉取
        info = self.get_table_info(table_name)
        upstream = info.get("dataworks_tasks", [])
        if not upstream:
            print(f"表 '{table_name}' 没有上游任务")
            return []

        results = []
        for task in upstream:
            task_id = task.get("Id")
            task_name = task.get("Name")
            result = self._get_task_code(task_id)
            if result:
                results.append(result)
            else:
                print(f"[WARN] 任务 {task_name}({task_id}) 所有 AK 均无权限")

        # 写缓存
        if results:
            parts = []
            for r in results:
                header = f"-- Task: {r['task_name']}  ID: {r['task_id']}  Type: {r['task_type']}"
                parts.append(f"{header}\n{r['content']}")
            self._write_cache(table_name, "\n\n".join(parts))
            print(f"[saved] {self._cache_path(table_name)}")

        # 获取并缓存 schema（复用已有的上游任务信息，避免重复 API 调用）
        up_tasks = [
            {"id": task.get("Id"), "name": task.get("Name")}
            for task in upstream
        ]
        self._ensure_schema_cache(table_name, force=force, dataworks_tasks=up_tasks)

        return results

    def get_upstream_tables(self, table_name: str) -> list[str]:
        """通过血缘 API 获取表的直接上游表列表。

        Returns:
            list of str，如 ["loghubods.user_share_log_flow", ...]
        """
        entity_id = self._build_entity_id(table_name)
        resp = _call_with_retry(lambda: self.client.list_lineages(
            _dw_models.ListLineagesRequest(dst_entity_id=entity_id, page_size=50)
        ))
        lineages = resp.body.to_map().get("PagingInfo", {}).get("Lineages", [])

        tables = []
        for l in lineages:
            src_id = l.get("SrcEntity", {}).get("Id", "")
            # maxcompute-table:::project::table → project.table
            parts = src_id.replace("maxcompute-table:::", "").split("::")
            if len(parts) == 2:
                tables.append(f"{parts[0]}.{parts[1]}")
        return sorted(set(tables))

    def get_node_code_recursive(self, table_name: str, max_depth: int = 3,
                                force: bool = False) -> dict:
        """BFS 逐层获取表及其所有上游表的生产代码。

        通过血缘 API（ListLineages）逐层追溯上游依赖，
        每层的代码和上游表都会被缓存到 production_code/。

        Args:
            table_name: 表名（支持 project.table 格式）
            max_depth: 最大追溯层数，默认 3
            force: True 时跳过缓存

        Returns:
            dict: {
                "project.table": {
                    "code": [...],            # get_node_code 返回值
                    "upstream": ["a.b", ...],  # 上游表名列表
                    "depth": int
                }, ...
            }
        """
        from collections import deque

        table_name = self._normalize_table_name(table_name)
        result = {}
        queue = deque([(table_name, 0)])
        visited = {table_name}

        while queue:
            tbl, depth = queue.popleft()
            indent = "  " * depth
            print(f"{indent}[depth={depth}] {tbl}")

            # 获取代码
            code = self.get_node_code(tbl, force=force)

            # 获取上游表
            upstream = []
            if depth < max_depth:
                try:
                    upstream = self.get_upstream_tables(tbl)
                except Exception:
                    pass

            result[tbl] = {"code": code, "upstream": upstream, "depth": depth}

            # 下一层入队
            for up_tbl in upstream:
                if up_tbl not in visited:
                    visited.add(up_tbl)
                    queue.append((up_tbl, depth + 1))

        # 打印汇总
        print(f"\n共追溯 {len(result)} 张表:")
        for tbl, info in result.items():
            has_code = "有代码" if info["code"] else "无代码"
            n_up = len(info["upstream"])
            print(f"  {'  ' * info['depth']}{tbl}  ({has_code}, {n_up} 个上游)")

        return result

    def print_node_code(self, table_name: str):
        """打印表的生产代码（人类可读格式）"""
        results = self.get_node_code(table_name)
        if not results:
            print(f"未找到 '{table_name}' 的生产代码")
            return
        for r in results:
            print(f"\n{'='*60}")
            print(f"任务: {r['task_name']}  ID: {r['task_id']}  "
                  f"类型: {r['task_type']}")
            print(f"{'='*60}")
            print(r["content"] or "(无内容)")
