"""dbt integration helpers for data modeling."""
from __future__ import annotations

from pathlib import Path
from typing import Dict, List, Optional
import yaml


def generate_dbt_assets(
    project_dir: Path,
    dictionary: Dict[str, Dict[str, Dict[str, str]]],
    relationships: List[dict],
    source_name: str = "raw",
    source_schema: Optional[str] = None,
) -> None:
    """Generate sources.yml and staging models from dictionary + relationships."""
    models_dir = project_dir / "models" / "staging"
    models_dir.mkdir(parents=True, exist_ok=True)

    source_block = {
        "name": source_name,
        "tables": [{"name": t.lower()} for t in dictionary.keys()],
    }
    if source_schema:
        source_block["schema"] = source_schema
    sources = {
        "version": 2,
        "sources": [
            source_block
        ],
    }

    (project_dir / "models" / "sources.yml").write_text(
        yaml.dump(sources, sort_keys=False),
        encoding="utf-8",
    )

    schema_models = {"version": 2, "models": []}

    for table, cols in dictionary.items():
        model_name = f"stg_{table.lower()}"
        sql = f"select * from {{{{ source('{source_name}','{table.lower()}') }}}}\n"
        (models_dir / f"{model_name}.sql").write_text(sql, encoding="utf-8")

        col_tests = []
        for rel in relationships:
            if rel["source_table"].lower() == table.lower():
                col_tests.append({
                    "name": rel["source_column"].lower(),
                    "tests": [
                        {
                            "relationships": {
                                "to": f"{{{{ ref('{ref_model_name(rel['target_table'])}') }}}}",
                                "field": rel["target_column"].lower(),
                            }
                        }
                    ],
                })

        schema_models["models"].append({
            "name": model_name,
            "description": f"Staging model for {table}.",
            "columns": col_tests,
        })

    (models_dir / "schema.yml").write_text(
        yaml.dump(schema_models, sort_keys=False),
        encoding="utf-8",
    )


def ref_model_name(table: str) -> str:
    return f"stg_{table.lower()}"


def ingest_dbt_artifacts(project_dir: Path) -> Dict[str, int]:
    """Ingest dbt manifest/catalog/run_results into KB stores."""
    from src.knowledge_base import KBIngestionPipeline
    from src.knowledge_base.types import GraphNode, GraphEdge, VectorMetadata

    pipeline = KBIngestionPipeline()
    target_dir = project_dir / "target"
    manifest_path = target_dir / "manifest.json"
    catalog_path = target_dir / "catalog.json"
    run_results_path = target_dir / "run_results.json"

    counts = {"manifest": 0, "catalog": 0, "run_results": 0}

    if manifest_path.exists():
        from src.data_catalog.lineage_extractor import DbtLineageExtractor

        extractor = DbtLineageExtractor()
        extraction = extractor.extract_from_manifest(str(manifest_path))
        nodes: List[GraphNode] = []
        edges: List[GraphEdge] = []
        vectors: List[VectorMetadata] = []

        for node in extraction.get("nodes", []):
            unique_id = node.get("unique_id", node.get("name"))
            node_id = f"DBT:{unique_id}"
            nodes.append(GraphNode(id=node_id, type="dbt_model", properties=node))
            vectors.append(VectorMetadata(
                id=node_id,
                source_type="dbt_model",
                source_id=unique_id,
                tags=["dbt", node.get("resource_type", "model")],
                extra={"name": node.get("name"), "materialized": node.get("materialized")},
            ))

        for edge in extraction.get("edges", []):
            edges.append(GraphEdge(
                source=f"DBT:{edge.get('source')}",
                target=f"DBT:{edge.get('target')}",
                type="DEPENDS_ON",
            ))

        if nodes:
            pipeline.ingest_graph(nodes, edges)
            pipeline.ingest_vector_metadata(vectors)
            counts["manifest"] = len(nodes)

    if catalog_path.exists():
        import json

        catalog = json.loads(catalog_path.read_text(encoding="utf-8"))
        nodes = catalog.get("nodes", {}) if isinstance(catalog, dict) else {}
        sources = catalog.get("sources", {}) if isinstance(catalog, dict) else {}
        summary = {
            "models": sum(1 for n in nodes.values() if n.get("resource_type") == "model"),
            "seeds": sum(1 for n in nodes.values() if n.get("resource_type") == "seed"),
            "sources": len(sources),
            "nodes_total": len(nodes),
        }
        pipeline.ingest_vector_metadata([
            VectorMetadata(
                id="DBT_CATALOG:snowflake_data_model",
                source_type="dbt_catalog",
                source_id="snowflake_data_model",
                tags=["dbt", "catalog"],
                extra={"summary": summary, "path": str(catalog_path)},
            )
        ])
        counts["catalog"] = summary["nodes_total"]

    if run_results_path.exists():
        import json

        run_results = json.loads(run_results_path.read_text(encoding="utf-8"))
        pipeline.ingest_vector_metadata([
            VectorMetadata(
                id="DBT_RUN_RESULTS:snowflake_data_model",
                source_type="dbt_run_results",
                source_id="snowflake_data_model",
                tags=["dbt", "run_results"],
                extra={"summary": run_results.get("metadata", {}), "path": str(run_results_path)},
            )
        ])
        counts["run_results"] = len(run_results.get("results", []))

    return counts
