"""Code generation workflow.

Reuses:
    - ``databridge.codegen.generate_code(query, output_type, output_dir)``
    - ``databridge.codegen.list_artifacts(output_dir)``

IKE integration (optional):
    - P2 Fix Engine: KB-grounded fix generation with trust FSM
    - P1 Knowledge Base: Query KB for context before generating
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

from databridge_core.sdk.core.context import SDKContext
from databridge_core.sdk.core.exceptions import NotImplementedStubError
from databridge_core.sdk.core.types import CompiledArtifact, Result
from databridge_core.sdk.workflows.base import BaseWorkflow
from databridge_core.sdk.workflows.result import WorkflowResult
from databridge_core.sdk.workflows.step import workflow_step


class CodeGenWorkflow(BaseWorkflow):
    """Generate code (SQL, mart DDL, Wright pipelines) from natural language.

    Uses a LangGraph pipeline: QueryAnalyzer -> ContextAssembler ->
    CodeGenerator -> Validator. Falls back to deterministic generation
    if LLM is unavailable.

    When IKE is available, the fix engine provides KB-grounded code
    generation using the 3-tier strategy (trust replay -> RAG ->
    LLM fallback) with graduated auto-approval via the Trust FSM.

    Reuses:
        - ``databridge.codegen.generate_code(query, output_type, output_dir)``
        - ``databridge.codegen.list_artifacts(output_dir)``
        - ``databridge.codegen.get_generated_artifact(artifact_id)``
    """

    name = "fix_generation.codegen"
    category = "fix_generation"
    description = "Generate SQL/DDL/pipeline code from natural language"
    required_params = ["query"]
    optional_params = {
        "output_type": "auto",
        "output_dir": "data/codegen",
        "ike_fix_engine": True,
    }

    def validate(self, params: Dict[str, Any]) -> Dict[str, Any]:
        self._check_required(params)
        return self._apply_defaults(params)

    def execute(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> Result:
        steps = []
        # IKE P1: Query KB for relevant context before generating
        if params.get("ike_fix_engine") and ctx and ctx.ike.available:
            steps.append(self._ike_query_context(params, ctx))
        steps.append(self._generate(params, ctx))
        # IKE P2: Generate fix using IKE's 3-tier engine
        if params.get("ike_fix_engine") and ctx and ctx.ike.available:
            steps.append(self._ike_generate_fix(params, ctx))
        # Graph integration
        if ctx and ctx.graphs.available:
            steps.append(self._graph_record_lineage(params, ctx))
            steps.append(self._graph_record_evidence(params, ctx))
            steps.append(self._graph_index_vectors(params, ctx))
        return WorkflowResult.from_steps(self.name, steps)

    def compile_to_snowflake(
        self, params: Dict[str, Any], ctx: Optional[SDKContext] = None
    ) -> List[CompiledArtifact]:
        raise NotImplementedStubError(
            self.name,
            "compile_to_snowflake",
            "The codegen output IS Snowflake DDL when output_type='sql' or 'mart'; "
            "wrap the generated SQL as CompiledArtifact directly",
        )

    @workflow_step("generate", "Generate code from query")
    def _generate(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # REUSE: await databridge.codegen.generate_code(
        #     query=params["query"],
        #     output_type=params["output_type"],
        #     output_dir=params["output_dir"],
        # )
        # NOTE: generate_code is async — needs asyncio.run() or event loop
        raise NotImplementedStubError(
            self.name,
            "generate",
            "databridge.codegen.generate_code(query, output_type, output_dir) — async",
        )

    @workflow_step("ike_query_context", "Query IKE KB for generation context")
    def _ike_query_context(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # IKE P1: Query the knowledge base for relevant context that
        # can inform code generation (prior fixes, patterns, rules).
        #
        # ctx.ike.query_knowledge(
        #     query=params["query"],
        #     top_k=5,
        # )
        raise NotImplementedStubError(
            self.name,
            "ike_query_context",
            "ctx.ike.query_knowledge(query, top_k=5) — IKE P1 hybrid KB search",
        )

    @workflow_step("ike_generate_fix", "Generate fix via IKE P2 engine")
    def _ike_generate_fix(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # IKE P2: Use the 3-tier fix engine for KB-grounded generation:
        #   Tier 1: Trust replay — reuse previously approved fix
        #   Tier 2: RAG — retrieve similar fixes from KB
        #   Tier 3: LLM fallback — generate new fix with LLM
        #
        # ctx.ike.generate_fix(
        #     rule_id="codegen",
        #     finding={"query": params["query"], "output_type": params["output_type"]},
        #     context=kb_context_from_prior_step,
        # )
        raise NotImplementedStubError(
            self.name,
            "ike_generate_fix",
            "ctx.ike.generate_fix(rule_id, finding, context) — IKE P2 3-tier fix engine",
        )

    # --- Graph integration ---

    @workflow_step("graph_record_lineage", "Record EXPORT_ARTIFACT node in lineage")
    def _graph_record_lineage(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.record_lineage_node("sdk", f"codegen:{artifact_id}", "EXTERNAL",
        #     description=f"Generated {params['output_type']} from: {params['query'][:50]}")
        raise NotImplementedStubError(
            self.name,
            "graph_record_lineage",
            "ctx.graphs.record_lineage_node — EXPORT_ARTIFACT for generated code",
        )

    @workflow_step("graph_record_evidence", "Record fix generation evidence")
    def _graph_record_evidence(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.record_evidence(
        #     account_id=params["query"][:50], hierarchy_node_id=artifact_id,
        #     method="codegen", confidence=generation_confidence,
        #     evidence_sources=[{"content": params["query"], "source_type": "user_query"}],
        # )
        raise NotImplementedStubError(
            self.name,
            "graph_record_evidence",
            "ctx.graphs.record_evidence — fix generation evidence trail",
        )

    @workflow_step("graph_index_vectors", "Index generated code in vector store")
    def _graph_index_vectors(self, params: Dict[str, Any], ctx: Optional[SDKContext] = None):
        # ctx.graphs.index_content(
        #     doc_id=f"codegen_{artifact_id}", content=generated_code,
        #     source_type="generated_code", metadata={"query": params["query"]},
        # )
        raise NotImplementedStubError(
            self.name,
            "graph_index_vectors",
            "ctx.graphs.index_content(doc_id, content, 'generated_code', metadata)",
        )
