"""
dbt-core adapter implementation.

Handles parsing dbt manifest.json/catalog.json and generating dbt schema YAML files.
"""

import json
import os
import re
import yaml
from pathlib import Path
from typing import Any, Optional

from trellis_datamodel.utils.yaml_handler import YamlHandler
from .base import (
    ColumnInfo,
    ColumnSchema,
    ModelInfo,
    ModelSchema,
    Relationship,
)


class DbtCoreAdapter:
    """Adapter for dbt-core transformation framework."""

    def __init__(
        self,
        manifest_path: str,
        catalog_path: str,
        project_path: str,
        data_model_path: str,
        model_paths: list[str],
    ):
        self.manifest_path = manifest_path
        self.catalog_path = catalog_path
        self.project_path = project_path
        self.data_model_path = data_model_path
        self.model_paths = model_paths
        self.yaml_handler = YamlHandler()

    def _load_catalog(self) -> Optional[dict]:
        """Load catalog.json if it exists."""
        if not os.path.exists(self.catalog_path):
            return None
        try:
            with open(self.catalog_path, "r") as f:
                return json.load(f)
        except Exception as exc:
            print(f"Warning: failed to read catalog at {self.catalog_path}: {exc}")
            return None

    def _load_manifest(self) -> dict:
        """Load manifest.json."""
        with open(self.manifest_path, "r") as f:
            return json.load(f)

    def _load_data_model(self) -> dict:
        """Load data model YAML if it exists."""
        if not self.data_model_path or not os.path.exists(self.data_model_path):
            return {}
        try:
            with open(self.data_model_path, "r") as f:
                return yaml.safe_load(f) or {}
        except Exception as e:
            print(f"Warning: Could not load data model: {e}")
            return {}

    def get_model_dirs(self) -> list[str]:
        """
        Resolve configured models directories, normalizing common dbt prefixes.

        Users may configure entries like "3_core", "models/3_entity", or absolute
        paths. We normalize these to real directories so downstream scans work.
        """

        def _normalize(subdir: str) -> str:
            # Absolute path - return as-is
            if os.path.isabs(subdir):
                return os.path.abspath(subdir).rstrip(os.sep)

            # Remove leading "./"
            while subdir.startswith("./"):
                subdir = subdir[2:]

            # Strip an optional leading "models/" so we don't double-prepend
            prefix = f"models{os.sep}"
            if subdir.startswith(prefix):
                subdir = subdir[len(prefix) :]

            return os.path.abspath(
                os.path.join(self.project_path, "models", subdir)
            ).rstrip(os.sep)

        if self.model_paths:
            # Remove duplicates while preserving order
            seen = set()
            normalized = []
            for path in self.model_paths:
                norm = _normalize(path)
                if norm not in seen:
                    seen.add(norm)
                    normalized.append(norm)
            return normalized

        return [
            os.path.abspath(os.path.join(self.project_path, "models")).rstrip(os.sep)
        ]

    def _entity_to_model_name(self, entity: dict[str, Any]) -> str:
        """
        Resolve the dbt model name for an entity.

        Prefers the bound dbt_model (strips project prefix), otherwise falls back to
        the entity ID so unbound entities still persist somewhere.
        """
        dbt_model = entity.get("dbt_model")
        if dbt_model:
            # dbt unique_id for versioned models looks like model.<project>.<name>.v2
            parts = dbt_model.split(".")
            if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
                # Use the model name part (the element before the vN suffix)
                return parts[-2]
            return parts[-1]
        return entity.get("id") or ""

    def _build_model_keys(self, base: str, version: Optional[str] = None) -> list[str]:
        """
        Generate a set of lookup keys for a model, including versioned variants.
        """
        keys = [base]
        if version is not None:
            # Version may come through as int from YAML; normalize to string
            version_str = str(version)
            # Support common ref patterns:
            # - ref('model', v=2)   -> base.v2
            # - alias names         -> base_v2
            # - fully qualified     -> base.v2 and base_v2
            version_num = version_str.lstrip("v")
            keys.extend(
                [
                    f"{base}.v{version_num}",
                    f"{base}_v{version_num}",
                    f"{base}v{version_num}",
                ]
            )
        # Deduplicate while preserving order
        seen = set()
        ordered_keys = []
        for k in keys:
            if k not in seen:
                seen.add(k)
                ordered_keys.append(k)
        return ordered_keys

    def _get_model_to_entity_map(self) -> dict[str, str]:
        """Build mapping from model names (with version aliases) to entity IDs."""
        model_to_entity: dict[str, str] = {}
        data_model = self._load_data_model()
        entities = data_model.get("entities", [])
        for entity in entities:
            entity_id = entity.get("id")
            dbt_model = entity.get("dbt_model")
            if dbt_model:
                parts = dbt_model.split(".")
                version_part = None
                if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
                    version_part = parts[-1].lstrip("v")
                    base_name = parts[-2]
                else:
                    base_name = parts[-1]

                # Map the raw unique_id as well as base/version variants
                model_to_entity[dbt_model] = entity_id
                for key in self._build_model_keys(base_name, version_part):
                    model_to_entity[key] = entity_id
            # Map additional models to the same entity
            additional_models = entity.get("additional_models", [])
            for add_model in additional_models:
                parts = add_model.split(".")
                version_part = None
                if len(parts) >= 2 and re.match(r"v\d+$", parts[-1]):
                    version_part = parts[-1].lstrip("v")
                    base_name = parts[-2]
                else:
                    base_name = parts[-1]

                model_to_entity[add_model] = entity_id
                for key in self._build_model_keys(base_name, version_part):
                    model_to_entity[key] = entity_id
            if entity_id:
                model_to_entity[entity_id] = entity_id
        return model_to_entity

    def _get_model_yml_path(
        self, model_name: str, target_version: Optional[int] = None
    ) -> Optional[str]:
        """Get the yml file path for a model from the manifest."""
        if not os.path.exists(self.manifest_path):
            return None

        manifest = self._load_manifest()
        preferred_node: Optional[dict] = None
        fallback_node: Optional[dict] = None
        for key, node in manifest.get("nodes", {}).items():
            if node.get("resource_type") == "model" and node.get("name") == model_name:
                if target_version is not None and node.get("version") == target_version:
                    preferred_node = node
                    break
                if fallback_node is None:
                    fallback_node = node

        node = preferred_node or fallback_node
        if not node:
            return None

        return self._derive_yml_path_from_node(node)

    def _normalize_patch_path(self, patch_path: str) -> str:
        """
        Convert a manifest patch_path (which may include a scheme) into an
        absolute filesystem path rooted at the dbt project.
        """
        if "://" in patch_path:
            patch_path = patch_path.split("://", 1)[1]

        patch_path = patch_path.lstrip("/")

        if os.path.isabs(patch_path):
            return patch_path

        # Default: resolve relative to the dbt project directory
        base = self.project_path or "."
        return os.path.abspath(os.path.join(base, patch_path))

    def _derive_yml_path_from_node(self, node: dict) -> Optional[str]:
        """
        Determine the YAML file path for a manifest node, preferring patch_path.
        """
        patch_path = node.get("patch_path")
        if patch_path:
            return self._normalize_patch_path(patch_path)

        original_file_path = node.get("original_file_path", "")
        if not original_file_path:
            return None

        sql_path = os.path.join(self.project_path, original_file_path)
        base_path = os.path.splitext(sql_path)[0]
        yml_path = f"{base_path}.yml"

        # Common versioned layout: player_v2.sql + player.yml
        if not os.path.exists(yml_path):
            base_path = re.sub(r"_v\d+$", "", base_path)
            yml_path = f"{base_path}.yml"

        return yml_path

    def _extract_version_from_string(self, value: str) -> Optional[int]:
        """Extract integer version from strings like 'model.proj.player.v2'."""
        if not value:
            return None

        match = re.search(r"\.v(\d+)$", value)
        if match:
            return int(match.group(1))

        if value.startswith("v") and value[1:].isdigit():
            return int(value[1:])

        return None

    def _resolve_model_version(
        self, model_name: str, entity_id: str, data_model: dict[str, Any]
    ) -> Optional[int]:
        """
        Resolve target version from data model (dbt_model) or manifest node.

        Prefers explicit dbt_model binding with .vN suffix, falls back to
        manifest node version.
        """
        for entity in data_model.get("entities", []):
            if entity.get("id") != entity_id:
                continue
            dbt_model = entity.get("dbt_model") or ""
            # Check last token after split to support fully-qualified model IDs
            last_token = dbt_model.split(".")[-1] if dbt_model else ""
            version = self._extract_version_from_string(last_token)
            if version is not None:
                return version

        if os.path.exists(self.manifest_path):
            manifest = self._load_manifest()
            for node in manifest.get("nodes", {}).values():
                if node.get("resource_type") != "model":
                    continue
                if node.get("name") != model_name:
                    continue
                node_version = node.get("version")
                if node_version:
                    return int(node_version)

        return None

    def _find_manifest_model_nodes(
        self, manifest: dict[str, Any], model_name: str
    ) -> list[dict[str, Any]]:
        """Collect manifest nodes matching a given model name."""
        return [
            node
            for node in manifest.get("nodes", {}).values()
            if node.get("resource_type") == "model" and node.get("name") == model_name
        ]

    def _select_model_node(
        self, candidates: list[dict[str, Any]], target_version: Optional[int]
    ) -> Optional[dict[str, Any]]:
        """
        Choose the best matching manifest node.

        Prefers an exact version match when requested; otherwise picks the
        highest numbered version when available, or the first candidate.
        """
        if target_version is not None:
            for node in candidates:
                node_version = node.get("version")
                if node_version is not None and int(node_version) == target_version:
                    return node

        if not candidates:
            return None

        versioned = [n for n in candidates if n.get("version") is not None]
        if versioned:
            return sorted(versioned, key=lambda n: n.get("version") or 0)[-1]

        return candidates[0]

    def get_models(self) -> list[ModelInfo]:
        """Parse dbt manifest and catalog to return available models."""
        if not os.path.exists(self.manifest_path):
            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")

        manifest = self._load_manifest()
        catalog = self._load_catalog()
        catalog_nodes = (catalog or {}).get("nodes", {})

        models: list[ModelInfo] = []
        for key, node in manifest.get("nodes", {}).items():
            if node.get("resource_type") != "model":
                continue

            # Filter by path
            original_path = node.get("original_file_path", "")
            if self.model_paths:
                match = any(pattern in original_path for pattern in self.model_paths)
                if not match:
                    continue

            # Get columns from catalog or manifest
            columns: list[ColumnInfo] = []
            unique_id = node.get("unique_id")
            catalog_node = catalog_nodes.get(unique_id)

            if catalog_node:
                for col in catalog_node.get("columns", {}).values():
                    columns.append(
                        {
                            "name": col.get("name"),
                            "type": col.get("type") or col.get("data_type"),
                        }
                    )
            else:
                for col_name, col_data in node.get("columns", {}).items():
                    columns.append({"name": col_name, "type": col_data.get("type")})

            # Extract materialization
            config = node.get("config", {})
            materialized = config.get("materialized", "view")

            models.append(
                {
                    "unique_id": unique_id,
                    "name": node.get("name"),
                    "version": node.get("version"),
                    "schema": node.get("schema"),
                    "table": node.get("alias", node.get("name")),
                    "columns": columns,
                    "description": node.get("description"),
                    "materialization": materialized,
                    "file_path": original_path,
                    "tags": node.get("tags", []),
                }
            )

        models.sort(key=lambda x: x["name"])
        return models

    def get_model_schema(
        self, model_name: str, version: Optional[int] = None
    ) -> ModelSchema:
        """Get the current schema definition for a specific model from its YAML file."""
        if not os.path.exists(self.manifest_path):
            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")

        manifest = self._load_manifest()

        candidate_nodes = self._find_manifest_model_nodes(manifest, model_name)
        model_node = self._select_model_node(candidate_nodes, target_version=version)

        if not model_node:
            raise ValueError(f"Model '{model_name}' not found in manifest")

        yml_path = self._derive_yml_path_from_node(model_node)
        if not yml_path:
            raise ValueError(
                f"No patch_path or original_file_path found for model '{model_name}'"
            )

        data = self.yaml_handler.load_file(yml_path)
        if not data:
            return {
                "model_name": model_name,
                "description": "",
                "columns": [],
                "tags": [],
                "file_path": yml_path,
            }

        model_entry = self.yaml_handler.find_model(data, model_name)
        if not model_entry:
            return {
                "model_name": model_name,
                "description": "",
                "columns": [],
                "tags": [],
                "file_path": yml_path,
            }

        target_version = version if version is not None else model_node.get("version")
        version_entry = None
        versions = model_entry.get("versions") or []
        if versions:
            if target_version is not None:
                for ver in versions:
                    if ver.get("v") == target_version:
                        version_entry = ver
                        break
            if version_entry is None:
                version_entry = sorted(versions, key=lambda ver: ver.get("v") or 0)[-1]

        # Prefer versioned block if present
        node_for_schema = version_entry or model_entry

        columns = self.yaml_handler.get_columns(node_for_schema)
        tags = self.yaml_handler.get_model_tags(node_for_schema)
        return {
            "model_name": model_name,
            "description": node_for_schema.get("description", ""),
            "columns": columns,
            "tags": tags,
            "file_path": yml_path,
        }

    def save_model_schema(
        self,
        model_name: str,
        columns: list[ColumnSchema],
        description: Optional[str] = None,
        tags: Optional[list[str]] = None,
        version: Optional[int] = None,
    ) -> Path:
        """Save/update the schema definition for a model."""
        if not os.path.exists(self.manifest_path):
            raise FileNotFoundError(f"Manifest not found at {self.manifest_path}")

        manifest = self._load_manifest()

        candidate_nodes = self._find_manifest_model_nodes(manifest, model_name)
        model_node = self._select_model_node(candidate_nodes, target_version=version)

        if not model_node:
            raise ValueError(f"Model '{model_name}' not found in manifest")

        yml_path = self._derive_yml_path_from_node(model_node)
        if not yml_path:
            raise ValueError(
                f"No patch_path or original_file_path found for model '{model_name}'"
            )

        data = self.yaml_handler.load_file(yml_path)
        if not data:
            data = {"version": 2, "models": []}

        model_entry = self.yaml_handler.ensure_model(data, model_name)

        target_version = version if version is not None else model_node.get("version")
        if target_version is not None:
            # Versioned model: update version entry and keep latest_version in sync (non-decreasing)
            self.yaml_handler.set_latest_version(model_entry, target_version)
            version_entry = self.yaml_handler.ensure_model_version(
                model_entry, target_version
            )

            if description is not None:
                self.yaml_handler.update_model_description(version_entry, description)
                # Keep top-level description aligned when present
                self.yaml_handler.update_model_description(model_entry, description)

            self.yaml_handler.update_columns_batch(version_entry, columns)

            if tags is not None:
                self.yaml_handler.update_version_tags(version_entry, tags)
        else:
            # Non-versioned model
            if description is not None:
                self.yaml_handler.update_model_description(model_entry, description)

            self.yaml_handler.update_columns_batch(model_entry, columns)

            if tags is not None:
                self.yaml_handler.update_model_tags(model_entry, tags)

        self.yaml_handler.save_file(yml_path, data)
        return Path(yml_path)

    def _parse_ref(self, ref_value: str) -> tuple[str, Optional[str]]:
        """
        Parse ref() targets, supporting optional version arguments.

        Examples:
            ref('player') -> ("player", None)
            ref('player', v=1) -> ("player", "1")
            ref("player", version=2) -> ("player", "2")
        """
        ref_pattern = (
            r"ref\(\s*['\"]([^,'\"]+)['\"](?:\s*,\s*(?:v|version)\s*=\s*([0-9]+))?\s*\)"
        )
        match = re.fullmatch(ref_pattern, ref_value.strip())
        if match:
            return match.group(1), match.group(2)
        return ref_value, None

    def _resolve_entity_id(
        self, model_to_entity: dict[str, str], base_name: str, version: Optional[str]
    ) -> str:
        """
        Resolve an entity id using base model name plus optional version.
        Falls back to the base name if no mapping is found.
        """
        for key in self._build_model_keys(base_name, version):
            if key in model_to_entity:
                return model_to_entity[key]
        # Fallback: try the raw base name
        return model_to_entity.get(base_name, base_name)

    def infer_relationships(self) -> list[Relationship]:
        """Scan dbt yml files and infer entity relationships from relationship tests."""
        model_dirs = self.get_model_dirs()
        model_to_entity = self._get_model_to_entity_map()

        # Only keep relationships where both ends map to entities that are bound to
        # at least one dbt model (including additional_models). This prevents writing
        # relationships for unbound entities in large projects.
        data_model = self._load_data_model()
        bound_entities = {
            e.get("id")
            for e in data_model.get("entities", [])
            if e.get("id") and (e.get("dbt_model") or e.get("additional_models"))
        }
        relationships: list[Relationship] = []
        yml_found = False

        for models_dir in model_dirs:
            if not os.path.exists(models_dir):
                continue

            for root, _, files in os.walk(models_dir):
                for filename in files:
                    if not filename.endswith((".yml", ".yaml")):
                        continue
                    yml_found = True

                    filepath = os.path.join(root, filename)
                    try:
                        with open(filepath, "r") as f:
                            schema_data = yaml.safe_load(f) or {}

                        models_list = schema_data.get("models", [])
                        for model in models_list:
                            base_model_name = model.get("name")
                            if not base_model_name:
                                continue

                            # Versioned models may declare columns inside versions list
                            version_entries = model.get("versions", [])
                            versioned_columns = []
                            if isinstance(version_entries, list) and version_entries:
                                for ver in version_entries:
                                    v_cols = ver.get("columns", [])
                                    versioned_columns.append(
                                        (
                                            ver.get("v") or ver.get("version"),
                                            v_cols,
                                        )
                                    )
                            else:
                                versioned_columns.append(
                                    (None, model.get("columns", []))
                                )

                            for model_version, columns in versioned_columns:
                                entity_id = self._resolve_entity_id(
                                    model_to_entity, base_model_name, model_version
                                )

                                for column in columns or []:
                                    test_blocks = []
                                    for key in ("tests", "data_tests"):
                                        value = column.get(key, [])
                                        if isinstance(value, list):
                                            test_blocks.extend(value)

                                    for test in test_blocks:
                                        if (
                                            not isinstance(test, dict)
                                            or "relationships" not in test
                                        ):
                                            continue

                                        rel_test = test["relationships"]
                                        args = rel_test.get("arguments", {}) or {}

                                        # Support both the recommended arguments block and legacy top-level keys
                                        to_ref = rel_test.get("to", "") or args.get(
                                            "to", ""
                                        )
                                        target_field = rel_test.get(
                                            "field", ""
                                        ) or args.get("field", "")

                                        # If either ref target or field is missing, skip and log for debugging
                                        if not to_ref or not target_field:
                                            continue

                                        target_base, target_version = self._parse_ref(
                                            to_ref
                                        )
                                        target_entity_id = self._resolve_entity_id(
                                            model_to_entity, target_base, target_version
                                        )

                                        # Skip relationships where either side is not bound to a
                                        # dbt model in the data model
                                        if (
                                            entity_id not in bound_entities
                                            or target_entity_id not in bound_entities
                                        ):
                                            continue

                                        relationships.append(
                                            {
                                                "source": target_entity_id,
                                                "target": entity_id,
                                                "label": "",
                                                "type": "one_to_many",
                                                "source_field": target_field,
                                                "target_field": column.get("name"),
                                            }
                                        )
                    except Exception as e:
                        print(f"Warning: Could not parse {filepath}: {e}")
                        continue

        if not yml_found:
            raise FileNotFoundError(
                f"No schema yml files found under configured dbt model paths: {model_dirs}"
            )

        # Remove duplicates
        seen: set[tuple] = set()
        unique_relationships: list[Relationship] = []
        for rel in relationships:
            key = (
                rel["source"],
                rel["target"],
                rel.get("source_field"),
                rel.get("target_field"),
            )
            if key not in seen:
                seen.add(key)
                unique_relationships.append(rel)

        return unique_relationships

    def sync_relationships(
        self,
        entities: list[dict[str, Any]],
        relationships: list[dict[str, Any]],
    ) -> list[Path]:
        """Sync relationship definitions from data model to dbt yml files."""
        # Build entity lookup
        entity_map = {e["id"]: e for e in entities if e.get("id")}
        # Map entity_id -> dbt model name (best-effort)
        entity_model_name: dict[str, str] = {
            eid: self._entity_to_model_name(ent) for eid, ent in entity_map.items()
        }

        # Group relationships by target entity (the one with the FK)
        fk_by_entity: dict[str, list[dict]] = {}

        for rel in relationships:
            source_id = rel.get("source")
            target_id = rel.get("target")
            rel_type = rel.get("type", "one_to_many")
            source_field = rel.get("source_field")
            target_field = rel.get("target_field")

            if not source_field or not target_field:
                continue

            fk_on_target = rel_type == "one_to_many"
            fk_entity = target_id if fk_on_target else source_id
            fk_field = target_field if fk_on_target else source_field
            ref_entity = source_id if fk_on_target else target_id
            ref_field = source_field if fk_on_target else target_field

            fk_by_entity.setdefault(fk_entity, []).append(
                {
                    "fk_field": fk_field,
                    "ref_entity": ref_entity,
                    "ref_field": ref_field,
                }
            )

        models_dir = self.get_model_dirs()[0]
        os.makedirs(models_dir, exist_ok=True)

        updated_files: list[Path] = []

        for entity in entities:
            entity_id = entity.get("id")
            if not entity_id:
                continue

            model_name = entity_model_name.get(entity_id, entity_id)

            # For bound entities, use the correct path from manifest
            # For unbound entities, fall back to models_dir/{entity_id}.yml
            yml_path = None
            if entity.get("dbt_model"):
                yml_path = self._get_model_yml_path(model_name)
            if not yml_path:
                yml_path = os.path.join(models_dir, f"{entity_id}.yml")

            data = self.yaml_handler.load_file(yml_path)
            if not data:
                data = {"version": 2, "models": []}

            model_entry = self.yaml_handler.ensure_model(data, model_name)

            if entity.get("description"):
                self.yaml_handler.update_model_description(
                    model_entry, entity.get("description")
                )

            # Sync Tags
            entity_tags = entity.get("tags")
            if entity_tags is not None:
                self.yaml_handler.update_model_tags(model_entry, entity_tags)

            # Sync Drafted Fields
            drafted_fields = entity.get("drafted_fields", [])
            for field in drafted_fields:
                f_name = field.get("name")
                f_type = field.get("datatype")
                f_desc = field.get("description")

                if not f_name:
                    continue

                col = self.yaml_handler.ensure_column(model_entry, f_name)
                self.yaml_handler.update_column(
                    col, data_type=f_type, description=f_desc
                )

            # Sync Relationships (FKs)
            fk_list = fk_by_entity.get(entity_id, [])
            for fk_info in fk_list:
                fk_field = fk_info["fk_field"]
                ref_entity = fk_info["ref_entity"]
                ref_field = fk_info["ref_field"]

                # Resolve reference model name (dbt model) for relationship test
                ref_model_name = entity_model_name.get(ref_entity, ref_entity)

                col = self.yaml_handler.ensure_column(model_entry, fk_field)

                if "data_type" not in col:
                    col["data_type"] = "text"

                self.yaml_handler.add_relationship_test(col, ref_model_name, ref_field)

            self.yaml_handler.save_file(yml_path, data)
            updated_files.append(Path(yml_path))

        return updated_files

    def save_dbt_schema(
        self,
        entity_id: str,
        model_name: str,
        fields: list[dict[str, str]],
        description: Optional[str] = None,
        tags: Optional[list[str]] = None,
    ) -> Path:
        """
        Generate and save a dbt schema YAML file for drafted fields.

        This is used for creating new schema files from the data model editor.
        """
        data_model = self._load_data_model()

        # Use description from request if available, otherwise fallback to data model
        entity_description = description
        if not entity_description:
            entities = data_model.get("entities", [])
            for entity in entities:
                if entity.get("id") == entity_id:
                    entity_description = entity.get("description")
                    break

        # Build a map of field names to relationships for this entity
        relationships = data_model.get("relationships", [])
        field_to_relationship: dict[str, dict] = {}
        # Map entity -> dbt model name for refs
        entity_model_name = {
            e.get("id"): self._entity_to_model_name(e)
            for e in data_model.get("entities", [])
            if e.get("id")
        }

        for rel in relationships:
            source_id = rel.get("source")
            target_id = rel.get("target")
            rel_type = rel.get("type", "one_to_many")
            source_field = rel.get("source_field")
            target_field = rel.get("target_field")

            if not source_field or not target_field:
                continue

            fk_on_target = rel_type == "one_to_many"
            fk_entity = target_id if fk_on_target else source_id
            fk_field = target_field if fk_on_target else source_field
            ref_entity = source_id if fk_on_target else target_id
            ref_field = source_field if fk_on_target else target_field

            if fk_entity == entity_id:
                field_to_relationship[fk_field] = {
                    "target_entity": ref_entity,
                    "target_field": ref_field,
                }

        # Generate YAML content with relationship tests
        columns = []
        for field in fields:
            column_dict: dict[str, Any] = {
                "name": field["name"],
                "data_type": field["datatype"],
            }

            if field.get("description"):
                column_dict["description"] = field["description"]

            field_name = field["name"]
            if field_name in field_to_relationship:
                rel_info = field_to_relationship[field_name]
                ref_model = entity_model_name.get(
                    rel_info["target_entity"], rel_info["target_entity"]
                )
                column_dict["data_tests"] = [
                    {
                        "relationships": {
                            "arguments": {
                                "to": f"ref('{ref_model}')",
                                "field": rel_info["target_field"],
                            }
                        }
                    }
                ]

            columns.append(column_dict)

        target_version = self._resolve_model_version(
            model_name=model_name, entity_id=entity_id, data_model=data_model
        )

        # Prefer manifest-derived path; fall back to default models directory
        yml_path = self._get_model_yml_path(model_name, target_version=target_version)
        if not yml_path:
            models_dir = self.get_model_dirs()[0]
            os.makedirs(models_dir, exist_ok=True)
            yml_path = os.path.join(models_dir, f"{entity_id}.yml")

        data = self.yaml_handler.load_file(yml_path)
        if not data:
            data = {"version": 2, "models": []}

        model_entry = self.yaml_handler.ensure_model(data, model_name)

        if target_version is not None:
            # Versioned model: update version block and latest_version
            self.yaml_handler.set_latest_version(model_entry, target_version)
            if entity_description:
                self.yaml_handler.update_model_description(
                    model_entry, entity_description
                )

            version_entry = self.yaml_handler.ensure_model_version(
                model_entry, target_version
            )
            self.yaml_handler.update_columns_batch(version_entry, columns)

            if entity_description:
                self.yaml_handler.update_model_description(
                    version_entry, entity_description
                )
            if tags is not None:
                self.yaml_handler.update_version_tags(version_entry, tags)
        else:
            # Non-versioned model: update columns and metadata directly
            if entity_description:
                self.yaml_handler.update_model_description(
                    model_entry, entity_description
                )

            self.yaml_handler.update_columns_batch(model_entry, columns)

            if tags is not None:
                self.yaml_handler.update_model_tags(model_entry, tags)

        self.yaml_handler.save_file(yml_path, data)
        return Path(yml_path)
