# User's semantic query (agnostic to physical storage)
query = """
    Find all customers with credit scores below 600 who had
    large transactions in the last 30 days,
    along with their associated risk alerts
"""

# coomia-dip needs to translate this to:
# 1. Query Iceberg table customer_profiles for credit scores
# 2. Query Iceberg table transactions for recent 30-day transactions
# 3. Query gRPC service for risk alerts
# 4. Inject tenant isolation conditions
# 5. Inject data masking rules
# 6. JOIN results and sort

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

@dataclass
class QueryContext:
    tenant_id: str
    world_id: str
    actor_id: str
    actor_roles: list[str]
    branch: str = "main"
    timestamp: str | None = None
    explain: bool = False
    metadata: dict[str, Any] = field(default_factory=dict)

@dataclass
class SemanticQuery:
    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:
    steps: list["PhysicalStep"]
    estimated_cost: float
    estimated_rows: int

@dataclass
class PhysicalStep:
    operation: str
    target: str
    details: dict[str, Any] = field(default_factory=dict)

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

    async def rewrite(self, raw_query: str, context: QueryContext) -> PhysicalPlan:
        semantic = await self._stages[0].process(raw_query, context)
        for stage in self._stages[1:]:
            semantic = await stage.process(semantic, context)
        return semantic

class RewriteStage(ABC):
    @abstractmethod
    async def process(self, input_data: Any, context: QueryContext) -> Any: ...

class QueryParser(RewriteStage):
    async def process(self, raw_query: str, context: QueryContext) -> SemanticQuery:
        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),
        )

class SemanticValidator(RewriteStage):
    async def process(self, query: SemanticQuery, context: QueryContext) -> SemanticQuery:
        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}'"
                        )

        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):
    async def process(self, query: SemanticQuery, context: QueryContext) -> SemanticQuery:
        # Tenant isolation
        query.filters.append({
            "field": "_tenant_id", "operator": "eq",
            "value": context.tenant_id, "injected": True,
        })

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

        # RBAC row-level filters
        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({
                    **row_filter, "injected": True, "source": "rbac",
                })

        # Remove inaccessible fields
        accessible = 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 or p == "*"
        ]
        return query

class PhysicalQueryGenerator(RewriteStage):
    async def process(self, query: SemanticQuery, context: QueryContext) -> PhysicalPlan:
        steps = []

        for obj_type in query.object_types:
            storage = await self._storage_registry.get_storage(obj_type, context.world_id)

            if storage.type == "iceberg":
                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":
                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}.")],
                    },
                ))

        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:
        projections = [
            p.split(".", 1)[1] for p in query.projections
            if p.startswith(f"{obj_type}.")
        ]
        proj_clause = ", ".join(projections) if projections else "*"

        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"

        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.limit:
            sql += f" LIMIT {query.limit}"
        if query.offset:
            sql += f" OFFSET {query.offset}"
        return sql

class PredicatePushdown(RewriteStage):
    async def process(self, plan: PhysicalPlan, context: QueryContext) -> PhysicalPlan:
        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 list(join_filters):
                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):
    async def process(self, plan: PhysicalPlan, context: QueryContext) -> PhysicalPlan:
        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"])

        for step in plan.steps:
            if step.operation == "scan":
                current = step.details.get("projections", [])
                if current and current != ["*"]:
                    step.details["projections"] = [
                        p for p in current if p in needed_fields or p in current
                    ]
        return plan

class DataMaskingInjector(RewriteStage):
    async def process(self, plan: PhysicalPlan, context: QueryContext) -> PhysicalPlan:
        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:
    async def rewrite_for_time_travel(
        self, query: SemanticQuery, timestamp: datetime, context: QueryContext
    ) -> SemanticQuery:
        context.timestamp = timestamp.isoformat()
        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:
        context.branch = branch
        return query

    async def diff_query(
        self, query: SemanticQuery, from_ref: str, to_ref: str, context: QueryContext
    ) -> dict:
        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:
    async def get_or_execute(
        self, query: SemanticQuery, context: QueryContext,
        executor: callable, ttl_seconds: int = 300
    ) -> list[dict]:
        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:
        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:
    async def explain(self, raw_query: str, context: QueryContext) -> dict:
        context.explain = True
        plan = await self._pipeline.rewrite(raw_query, context)
        return {
            "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")
            ],
        }