from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field, field_validator, model_validator


class CodeChunkDocument(BaseModel):
    """Upstream `code_chunks` document (Code Parser)."""

    chunk_id: str
    repo_id: str
    snapshot_id: str
    commit_sha: str
    file_path: str
    symbol_name: str
    text: str
    start_line: int
    end_line: int
    graph_node_id: str | None = None
    language: str | None = None
    symbol_type: str | None = None
    chunk_hash: str | None = None
    metadata: dict[str, Any] = Field(default_factory=dict)

    @field_validator("metadata", mode="before")
    @classmethod
    def _default_metadata(cls, value: Any) -> dict[str, Any]:
        return value or {}

    @model_validator(mode="before")
    @classmethod
    def _coalesce_aliases(cls, data: Any) -> Any:
        if not isinstance(data, dict):
            return data

        merged = dict(data)
        if not merged.get("text"):
            for key in ("chunk_text", "symbol_text"):
                if merged.get(key):
                    merged["text"] = merged[key]
                    break

        if merged.get("chunk_hash") is None and isinstance(merged.get("metadata"), dict):
            merged["chunk_hash"] = merged["metadata"].get("chunk_hash")

        for field, meta_key in (
            ("language", "language"),
            ("symbol_type", "symbol_type"),
        ):
            if merged.get(field) is None and isinstance(merged.get("metadata"), dict):
                merged[field] = merged["metadata"].get(meta_key)

        return merged