"""MetricFlow SDK client and query result data model."""

from __future__ import annotations

import json
import logging
import threading
from dataclasses import dataclass
from pathlib import Path
import datetime
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union

from dbt_semantic_interfaces.protocols import SemanticManifest

from metricflow.engine.metricflow_engine import MetricFlowQueryRequest, MetricFlowQueryResult
from metricflow.engine.models import Dimension, Entity, Metric, SavedQuery
from metricflow.engine.sdk_exceptions import ClientClosedError, InitializationError, QueryError
from metricflow.protocols.sql_client import SqlClient

__all__ = [
    "MetricFlowClient",
    "QueryResult",
]


@dataclass(frozen=True)
class QueryResult:
    """SDK 查询结果。"""

    column_names: Tuple[str, ...]
    rows: Tuple[Tuple[Any, ...], ...]
    sql: str

    @property
    def row_count(self) -> int:
        """返回结果行数。"""
        return len(self.rows)

    def to_dict(self) -> Dict[str, Any]:
        """转换为字典格式。"""
        return {
            "column_names": list(self.column_names),
            "rows": [list(row) for row in self.rows],
            "sql": self.sql,
        }

    def to_json(self) -> str:
        """转换为 JSON 字符串。"""
        return json.dumps(self.to_dict(), default=str)


class MetricFlowClient:
    """MetricFlow SDK 的主入口，提供编程式 API。"""

    def __init__(
        self,
        semantic_manifest: SemanticManifest,
        sql_client: SqlClient,
    ) -> None:
        """从已有的 SemanticManifest 和 SqlClient 初始化。

        Args:
            semantic_manifest: dbt 语义清单对象。
            sql_client: 数据库连接客户端。

        Raises:
            InitializationError: 初始化过程中发生任何错误时抛出。
        """
        try:
            from metricflow.engine.metricflow_engine import MetricFlowEngine
            from metricflow_semantics.model.semantic_manifest_lookup import SemanticManifestLookup

            lookup = SemanticManifestLookup(semantic_manifest)
            self._engine = MetricFlowEngine(semantic_manifest_lookup=lookup, sql_client=sql_client)
        except Exception as e:
            raise InitializationError(f"MetricFlowClient 初始化失败: {e}") from e

        self._sql_client = sql_client
        self._closed = False
        self._lock = threading.Lock()

    @classmethod
    def from_dbt_project(
        cls,
        project_path: Union[str, Path],
        profiles_path: Optional[Union[str, Path]] = None,
    ) -> MetricFlowClient:
        """从 dbt project 路径初始化。需要安装 dbt-metricflow。

        Args:
            project_path: dbt 项目根目录路径。
            profiles_path: profiles.yml 所在目录路径，默认与 project_path 相同。

        Raises:
            InitializationError: dbt-metricflow 未安装或加载失败时抛出。
        """
        try:
            from dbt_metricflow.cli.dbt_connectors.adapter_backed_client import AdapterBackedSqlClient
            from dbt_metricflow.cli.dbt_connectors.dbt_config_accessor import dbtArtifacts, dbtProjectMetadata
        except ImportError as e:
            raise InitializationError(
                "从 dbt project 初始化需要安装 dbt-metricflow 包。" "请运行: pip install dbt-metricflow"
            ) from e

        try:
            resolved_project_path = Path(project_path)
            resolved_profiles_path = Path(profiles_path) if profiles_path is not None else resolved_project_path

            project_metadata = dbtProjectMetadata.load_from_paths(
                profiles_path=resolved_profiles_path,
                project_path=resolved_project_path,
            )
            artifacts = dbtArtifacts.load_from_project_metadata(project_metadata)
            sql_client = AdapterBackedSqlClient(artifacts.adapter)
            return cls(semantic_manifest=artifacts.semantic_manifest, sql_client=sql_client)
        except InitializationError:
            raise
        except Exception as e:
            raise InitializationError(f"从 dbt project 加载失败: {e}") from e

    def _check_not_closed(self) -> None:
        """检查 client 是否已关闭。"""
        with self._lock:
            if self._closed:
                raise ClientClosedError("MetricFlowClient 已关闭，无法执行操作。")

    def close(self) -> None:
        """关闭底层连接。幂等操作。"""
        with self._lock:
            if not self._closed:
                self._sql_client.close()
                self._closed = True

    def __enter__(self) -> MetricFlowClient:
        return self

    def __exit__(
        self,
        exc_type: Optional[type],
        exc_val: Optional[BaseException],
        exc_tb: Optional[Any],
    ) -> None:
        self.close()

    def __del__(self) -> None:
        if getattr(self, "_closed", True) is False:
            logging.getLogger(__name__).warning("MetricFlowClient 未被正确关闭。请使用 close() 或 context manager。")

    def _convert_result(self, mf_result: MetricFlowQueryResult) -> QueryResult:
        """将 MetricFlowQueryResult 转换为 SDK QueryResult。"""
        df = mf_result.result_df
        if df is None:
            return QueryResult(column_names=(), rows=(), sql=mf_result.sql)
        return QueryResult(
            column_names=tuple(df.column_names),
            rows=tuple(tuple(row) for row in df.rows),
            sql=mf_result.sql,
        )

    def query(
        self,
        metrics: Optional[Sequence[str]] = None,
        group_by: Optional[Sequence[str]] = None,
        where: Optional[Sequence[str]] = None,
        order_by: Optional[Sequence[str]] = None,
        limit: Optional[int] = None,
        time_constraint_start: Optional[datetime.datetime] = None,
        time_constraint_end: Optional[datetime.datetime] = None,
        saved_query: Optional[str] = None,
    ) -> QueryResult:
        """执行指标查询。

        Args:
            metrics: 要查询的指标名称列表。
            group_by: 分组维度名称列表。
            where: WHERE 条件列表。
            order_by: 排序字段名称列表。
            limit: 返回行数限制。
            time_constraint_start: 时间约束起始。
            time_constraint_end: 时间约束结束。
            saved_query: 已保存查询的名称。

        Returns:
            QueryResult: 包含列名、数据行和 SQL 的查询结果。

        Raises:
            QueryError: 查询执行失败时抛出。
        """
        self._check_not_closed()
        try:
            mf_request = MetricFlowQueryRequest.create_with_random_request_id(
                saved_query_name=saved_query,
                metric_names=metrics,
                group_by_names=group_by,
                where_constraints=where,
                order_by_names=order_by,
                limit=limit,
                time_constraint_start=time_constraint_start,
                time_constraint_end=time_constraint_end,
            )
            mf_result = self._engine.query(mf_request)
            return self._convert_result(mf_result)
        except Exception as e:
            raise QueryError(f"查询执行失败: {e}") from e

    def explain(
        self,
        metrics: Optional[Sequence[str]] = None,
        group_by: Optional[Sequence[str]] = None,
        where: Optional[Sequence[str]] = None,
        order_by: Optional[Sequence[str]] = None,
        limit: Optional[int] = None,
        time_constraint_start: Optional[datetime.datetime] = None,
        time_constraint_end: Optional[datetime.datetime] = None,
        saved_query: Optional[str] = None,
    ) -> str:
        """返回将要执行的 SQL，不实际执行查询。

        Raises:
            QueryError: 生成 SQL 失败时抛出。
        """
        self._check_not_closed()
        try:
            mf_request = MetricFlowQueryRequest.create_with_random_request_id(
                saved_query_name=saved_query,
                metric_names=metrics,
                group_by_names=group_by,
                where_constraints=where,
                order_by_names=order_by,
                limit=limit,
                time_constraint_start=time_constraint_start,
                time_constraint_end=time_constraint_end,
            )
            explain_result = self._engine.explain(mf_request)
            return explain_result.sql_statement.sql
        except Exception as e:
            raise QueryError(f"SQL 生成失败: {e}") from e

    def list_metrics(self) -> List[Metric]:
        """列出所有可用指标。

        Raises:
            QueryError: 查询失败时抛出。
        """
        self._check_not_closed()
        try:
            return self._engine.list_metrics()
        except Exception as e:
            raise QueryError(f"列出指标失败: {e}") from e

    def list_dimensions(self, metric_names: Sequence[str]) -> List[Dimension]:
        """列出指定指标的共有维度。

        Raises:
            QueryError: 查询失败时抛出。
        """
        self._check_not_closed()
        try:
            return self._engine.list_dimensions(metric_names=list(metric_names))
        except Exception as e:
            raise QueryError(f"列出维度失败: {e}") from e

    def list_entities(self, metric_names: Sequence[str]) -> List[Entity]:
        """列出指定指标的共有实体。

        Raises:
            QueryError: 查询失败时抛出。
        """
        self._check_not_closed()
        try:
            return self._engine.entities_for_metrics(metric_names=list(metric_names))
        except Exception as e:
            raise QueryError(f"列出实体失败: {e}") from e

    def list_saved_queries(self) -> List[SavedQuery]:
        """列出所有已保存查询。

        Raises:
            QueryError: 查询失败时抛出。
        """
        self._check_not_closed()
        try:
            return self._engine.list_saved_queries()
        except Exception as e:
            raise QueryError(f"列出已保存查询失败: {e}") from e

    def get_dimension_values(
        self,
        metric_names: Sequence[str],
        dimension: str,
        time_constraint_start: Optional[datetime.datetime] = None,
        time_constraint_end: Optional[datetime.datetime] = None,
    ) -> List[str]:
        """获取维度的所有可用值。

        Raises:
            QueryError: 查询失败时抛出。
        """
        self._check_not_closed()
        try:
            return self._engine.get_dimension_values(
                metric_names=list(metric_names),
                get_group_by_values=dimension,
                time_constraint_start=time_constraint_start,
                time_constraint_end=time_constraint_end,
            )
        except Exception as e:
            raise QueryError(f"获取维度值失败: {e}") from e
