# 用户的语义查询（不关心物理存储）
query = """
    找到所有信用评分低于 600 且最近 30 天有大额交易的客户，
    以及他们关联的风控告警
"""

# coomia-dip 需要将此翻译为：
# 1. 从 Iceberg 表 customer_profiles 查询信用评分
# 2. 从 Iceberg 表 transactions 查询近 30 天交易
# 3. 从 gRPC 服务查询风控告警
# 4. 注入租户隔离条件
# 5. 注入数据脱敏规则
# 6. JOIN 结果并排序

from dataclasses import dataclass, field
from typing import Any
from abc import ABC, abstractmethod

@dataclass
class QueryContext:
    """Context for query rewriting."""
    tenant_id: str
    world_id: str
    actor_id: str
    actor_roles: list[str]
    branch: str = "main"          # Nessie branch
    timestamp: str | None = None  # 时间旅行查询
    explain: bool = False         # 是否返回执行计划
    metadata: dict[str, Any] = field(default_factory=dict)

@dataclass
class SemanticQuery:
    """Parsed semantic query."""
    object_types: list[str]
    filters: list[dict]
    projections: list[str]
    joins: list[dict]
    order_by: list[dict]
    limit: int | None = None
    offset: int | None = None
    aggregations: list[dict] = field(default_factory=list)

@dataclass
class PhysicalPlan:
    """Physical execution plan."""
    steps: list["PhysicalStep"]
    estimated_cost: float
    estimated_rows: int

@dataclass
class PhysicalStep:
    """A step in the physical execution plan."""
    operation: str          # scan | filter | join | aggregate | sort
    target: str             # table name or service endpoint
    details: dict[str, Any] = field(default_factory=dict)

class QueryRewritePipeline:
    """Multi-stage query rewrite pipeline."""

    def __init__(self, stages: list["RewriteStage"]):
        self._stages = stages

    async def rewrite(
        self,
        raw_query: str,
        context: QueryContext,
    ) -> PhysicalPlan:
        """Rewrite a semantic query into a physical execution plan."""
        # Stage 1: Parse
        semantic = await self._stages[0].process(raw_query, context)

        # Stage 2-N: Apply each rewrite stage
        for stage in self._stages[1:]:
            semantic = await stage.process(semantic, context)

        return semantic

class RewriteStage(ABC):
    """Abstract base for query rewrite stages."""

    @abstractmethod
    async def process(self, input_data: Any, context: QueryContext) -> Any:
        """Process input and return transformed output."""
        ...

class QueryParser(RewriteStage):
    """Parse raw query into semantic representation."""

    async def process(self, raw_query: str, context: QueryContext) -> SemanticQuery:
        """Parse the ontology query language (OQL)."""
        tokens = self._tokenize(raw_query)
        ast = self._build_ast(tokens)

        return SemanticQuery(
            object_types=self._extract_types(ast),
            filters=self._extract_filters(ast),
            projections=self._extract_projections(ast),
            joins=self._extract_joins(ast),
            order_by=self._extract_order_by(ast),
            limit=self._extract_limit(ast),
            offset=self._extract_offset(ast),
            aggregations=self._extract_aggregations(ast),
        )

    def _tokenize(self, query: str) -> list:
        """Tokenize the query string."""
        # OQL 语法解析
        pass

    def _build_ast(self, tokens: list) -> dict:
        """Build abstract syntax tree."""
        pass

class SemanticValidator(RewriteStage):
    """Validate query against Ontology schema."""

    async def process(
        self, query: SemanticQuery, context: QueryContext
    ) -> SemanticQuery:
        """Validate that all referenced types and properties exist."""
        for obj_type in query.object_types:
            schema = await self._schema_registry.get(obj_type, context.world_id)
            if not schema:
                raise QueryValidationError(
                    f"ObjectType '{obj_type}' not found in world '{context.world_id}'"
                )

            # 验证字段存在
            for proj in query.projections:
                if proj.startswith(f"{obj_type}."):
                    field_name = proj.split(".", 1)[1]
                    if field_name not in schema["properties"]:
                        raise QueryValidationError(
                            f"Property '{field_name}' not found in '{obj_type}'"
                        )

        # 验证 JOIN 关系存在
        for join in query.joins:
            link = await self._schema_registry.get_link_type(
                join["from_type"], join["to_type"]
            )
            if not link:
                raise QueryValidationError(
                    f"No relationship between '{join['from_type']}' and '{join['to_type']}'"
                )

        return query

class PermissionInjector(RewriteStage):
    """Inject permission filters into the query."""

    async def process(
        self, query: SemanticQuery, context: QueryContext
    ) -> SemanticQuery:
        """Add tenant isolation and RBAC filters."""
        # 注入租户隔离
        query.filters.append({
            "field": "_tenant_id",
            "operator": "eq",
            "value": context.tenant_id,
            "injected": True,  # 标记为系统注入的过滤条件
        })

        # 注入 World 隔离
        query.filters.append({
            "field": "_world_id",
            "operator": "eq",
            "value": context.world_id,
            "injected": True,
        })

        # 注入 RBAC 行级过滤
        for obj_type in query.object_types:
            row_filter = await self._rbac_service.get_row_filter(
                context.actor_id, context.actor_roles, obj_type
            )
            if row_filter:
                query.filters.append({
                    "field": row_filter["field"],
                    "operator": row_filter["operator"],
                    "value": row_filter["value"],
                    "injected": True,
                    "source": "rbac",
                })

        # 移除无权限的字段
        accessible_fields = await self._rbac_service.get_accessible_fields(
            context.actor_id, context.actor_roles, query.object_types
        )
        query.projections = [
            p for p in query.projections
            if p in accessible_fields or p == "*"
        ]

        return query

class PhysicalQueryGenerator(RewriteStage):
    """Generate physical execution plan from semantic query."""

    async def process(
        self, query: SemanticQuery, context: QueryContext
    ) -> PhysicalPlan:
        """Convert semantic query to physical plan."""
        steps = []

        for obj_type in query.object_types:
            # 确定物理存储位置
            storage = await self._storage_registry.get_storage(
                obj_type, context.world_id
            )

            if storage.type == "iceberg":
                # 生成 Iceberg/Trino SQL
                sql = self._generate_iceberg_sql(query, obj_type, storage, context)
                steps.append(PhysicalStep(
                    operation="scan",
                    target=storage.table_name,
                    details={
                        "sql": sql,
                        "branch": context.branch,
                        "timestamp": context.timestamp,
                        "engine": "trino",
                    },
                ))
            elif storage.type == "grpc_service":
                # 生成 gRPC 调用
                steps.append(PhysicalStep(
                    operation="rpc_call",
                    target=storage.endpoint,
                    details={
                        "service": storage.service_name,
                        "method": "Query",
                        "filters": [f for f in query.filters if self._applies_to(f, obj_type)],
                        "projections": [p for p in query.projections if p.startswith(f"{obj_type}.")],
                    },
                ))

        # 添加 JOIN 步骤
        for join in query.joins:
            steps.append(PhysicalStep(
                operation="join",
                target="in_memory",
                details={
                    "type": join.get("type", "inner"),
                    "left": join["from_type"],
                    "right": join["to_type"],
                    "on": join["on"],
                },
            ))

        # 添加聚合步骤
        if query.aggregations:
            steps.append(PhysicalStep(
                operation="aggregate",
                target="in_memory",
                details={"aggregations": query.aggregations},
            ))

        # 添加排序和分页
        if query.order_by:
            steps.append(PhysicalStep(
                operation="sort",
                target="in_memory",
                details={"order_by": query.order_by},
            ))

        return PhysicalPlan(
            steps=steps,
            estimated_cost=self._estimate_cost(steps),
            estimated_rows=self._estimate_rows(steps),
        )

    def _generate_iceberg_sql(
        self,
        query: SemanticQuery,
        obj_type: str,
        storage: "StorageInfo",
        context: QueryContext,
    ) -> str:
        """Generate SQL for Iceberg table query."""
        projections = [
            p.split(".", 1)[1] for p in query.projections
            if p.startswith(f"{obj_type}.")
        ]
        proj_clause = ", ".join(projections) if projections else "*"

        # 构建 WHERE 子句
        where_parts = []
        for f in query.filters:
            if self._applies_to(f, obj_type):
                field = f["field"].replace(f"{obj_type}.", "")
                where_parts.append(
                    f"{field} {f['operator']} {self._format_value(f['value'])}"
                )

        where_clause = " AND ".join(where_parts) if where_parts else "1=1"

        # Nessie 分支和时间旅行
        table_ref = f"{storage.catalog}.{storage.schema_name}.{storage.table_name}"
        if context.branch != "main":
            table_ref += f"@{context.branch}"
        if context.timestamp:
            table_ref += f" FOR SYSTEM_TIME AS OF TIMESTAMP '{context.timestamp}'"

        sql = f"SELECT {proj_clause} FROM {table_ref} WHERE {where_clause}"

        if query.order_by:
            order_parts = [
                f"{o['field'].replace(f'{obj_type}.', '')} {o.get('direction', 'ASC')}"
                for o in query.order_by
                if o["field"].startswith(f"{obj_type}.")
            ]
            if order_parts:
                sql += f" ORDER BY {', '.join(order_parts)}"

        if query.limit:
            sql += f" LIMIT {query.limit}"
        if query.offset:
            sql += f" OFFSET {query.offset}"

        return sql

class PredicatePushdown(RewriteStage):
    """Push predicates closer to data sources."""

    async def process(
        self, plan: PhysicalPlan, context: QueryContext
    ) -> PhysicalPlan:
        """Move filters to earliest possible execution point."""
        join_steps = [s for s in plan.steps if s.operation == "join"]
        scan_steps = [s for s in plan.steps if s.operation in ("scan", "rpc_call")]

        for join_step in join_steps:
            join_filters = join_step.details.get("filters", [])
            for f in join_filters:
                # 如果过滤条件只涉及一个表，下推到该表的 scan 步骤
                target_type = self._get_filter_type(f)
                for scan in scan_steps:
                    if scan.target == target_type:
                        scan.details.setdefault("filters", []).append(f)
                        join_filters.remove(f)
                        break

        return plan

class ProjectionPruning(RewriteStage):
    """Remove unnecessary columns from scan operations."""

    async def process(
        self, plan: PhysicalPlan, context: QueryContext
    ) -> PhysicalPlan:
        """Only fetch columns that are actually needed."""
        # 收集所有需要的字段
        needed_fields: set[str] = set()
        for step in plan.steps:
            if step.operation == "join":
                needed_fields.update(step.details.get("on", {}).values())
            if step.operation == "sort":
                for o in step.details.get("order_by", []):
                    needed_fields.add(o["field"])
            if step.operation == "aggregate":
                for a in step.details.get("aggregations", []):
                    needed_fields.add(a["field"])

        # 裁剪 scan 步骤中不需要的列
        for step in plan.steps:
            if step.operation == "scan":
                current_projections = step.details.get("projections", [])
                if current_projections and current_projections != ["*"]:
                    step.details["projections"] = [
                        p for p in current_projections
                        if p in needed_fields or p in current_projections
                    ]

        return plan

class DataMaskingInjector(RewriteStage):
    """Inject data masking rules into query results."""

    async def process(
        self, plan: PhysicalPlan, context: QueryContext
    ) -> PhysicalPlan:
        """Add masking transformations to sensitive fields."""
        masking_rules = await self._masking_service.get_rules(
            context.actor_id, context.actor_roles
        )

        if masking_rules:
            plan.steps.append(PhysicalStep(
                operation="transform",
                target="in_memory",
                details={
                    "type": "masking",
                    "rules": masking_rules,
                },
            ))

        return plan

class TimeTravelQueryRewriter:
    """Rewrite queries for time-travel capabilities."""

    async def rewrite_for_time_travel(
        self,
        query: SemanticQuery,
        timestamp: datetime,
        context: QueryContext,
    ) -> SemanticQuery:
        """Rewrite query to access data at a specific point in time."""
        context.timestamp = timestamp.isoformat()

        # 获取该时间点的 Nessie commit
        commit = await self._nessie.get_commit_at(
            context.branch, timestamp
        )
        context.metadata["nessie_commit"] = commit.hash

        return query

    async def rewrite_for_branch(
        self,
        query: SemanticQuery,
        branch: str,
        context: QueryContext,
    ) -> SemanticQuery:
        """Rewrite query to access data on a specific branch."""
        context.branch = branch
        return query

    async def diff_query(
        self,
        query: SemanticQuery,
        from_ref: str,
        to_ref: str,
        context: QueryContext,
    ) -> dict:
        """Query the difference between two points in time."""
        results_from = await self._executor.execute(
            query, QueryContext(**{**vars(context), "branch": from_ref})
        )
        results_to = await self._executor.execute(
            query, QueryContext(**{**vars(context), "branch": to_ref})
        )

        return {
            "added": [r for r in results_to if r not in results_from],
            "removed": [r for r in results_from if r not in results_to],
            "modified": self._find_modified(results_from, results_to),
        }

class SemanticQueryCache:
    """Cache query results at the semantic level."""

    async def get_or_execute(
        self,
        query: SemanticQuery,
        context: QueryContext,
        executor: callable,
        ttl_seconds: int = 300,
    ) -> list[dict]:
        """Get cached results or execute and cache."""
        cache_key = self._compute_cache_key(query, context)

        # 检查缓存
        cached = await self._cache.get(cache_key)
        if cached is not None:
            return cached

        # 执行查询
        results = await executor(query, context)

        # 缓存结果
        await self._cache.set(cache_key, results, ttl=ttl_seconds)

        return results

    def _compute_cache_key(
        self, query: SemanticQuery, context: QueryContext
    ) -> str:
        """Compute a deterministic cache key."""
        import hashlib, json
        key_data = {
            "types": sorted(query.object_types),
            "filters": sorted(str(f) for f in query.filters),
            "projections": sorted(query.projections),
            "tenant": context.tenant_id,
            "world": context.world_id,
            "branch": context.branch,
            "timestamp": context.timestamp,
        }
        return hashlib.sha256(
            json.dumps(key_data, sort_keys=True).encode()
        ).hexdigest()

class QueryExplainer:
    """Explain query execution plan for debugging and optimization."""

    async def explain(
        self, raw_query: str, context: QueryContext
    ) -> dict:
        """Generate a detailed explanation of the query execution plan."""
        context.explain = True
        plan = await self._pipeline.rewrite(raw_query, context)

        explanation = {
            "query": raw_query,
            "plan": {
                "steps": [
                    {
                        "operation": s.operation,
                        "target": s.target,
                        "details": s.details,
                    }
                    for s in plan.steps
                ],
                "estimated_cost": plan.estimated_cost,
                "estimated_rows": plan.estimated_rows,
            },
            "optimizations_applied": [
                "predicate_pushdown",
                "projection_pruning",
                "data_masking_injection",
            ],
            "injected_filters": [
                f for step in plan.steps
                for f in step.details.get("filters", [])
                if f.get("injected")
            ],
        }

        return explanation