from dataclasses import dataclass

from vespa.package import (
    ApplicationPackage,
    Document,
    DocumentSummary,
    Field,
    FieldSet,
    Function,
    RankProfile,
    Schema,
    SecondPhaseRanking,
    Summary,
)


@dataclass
class VespaSchema:
    schema_name: str
    app_package_name: str
    embedding_dim: int = 384
    chunk_size: int = 1024
    distance_metric: str = "angular"

    def create_schema_fields(self) -> Schema:
        """Create a Vespa schema with predefined fields and rank profile."""
        # Create document with basic fields
        document = Document(
            fields=[
                Field(
                    name="id",
                    type="string",
                    indexing=["attribute", "summary"],
                ),
                Field(
                    name="content",
                    type="string",
                    indexing=["index", "summary"],
                    index="enable-bm25",
                ),
                Field(
                    name="loc",
                    type="string",
                    indexing=["index", "summary"],
                    index="enable-bm25",
                ),
                Field(
                    name="title",
                    type="string",
                    indexing=["index", "summary"],
                    index="enable-bm25",
                ),
                Field(
                    name="chunks",
                    type="array<string>",
                    indexing=["summary", "index"],
                    index="enable-bm25",
                ),
                Field(
                    name="chunk_count",
                    type="int",
                    indexing=["attribute", "summary"],
                ),
                Field(
                    name="content_embedding",
                    type=f"tensor<float>(x[{self.embedding_dim}])",
                    indexing=["attribute", "index"],
                    attribute=[f"distance-metric: {self.distance_metric}"],
                ),
                Field(
                    name="chunk_embeddings",
                    type=f"tensor<float>(chunk{{}}, x[{self.embedding_dim}])",
                    indexing=["attribute", "index"],
                    attribute=[f"distance-metric: {self.distance_metric}"],
                ),
            ]
        )

        # Create schema with document
        schema = Schema(name=self.schema_name, document=document)

        # Define field set
        schema.add_field_set(FieldSet(name="default", fields=["content", "chunks"]))

        # Add document summaries
        self._add_document_summaries(schema)

        return schema

    def _add_document_summaries(self, schema: Schema) -> None:
        """Add document summaries to the schema."""
        # Document summary: no-chunks (basic fields without chunks content)
        schema.add_document_summary(
            DocumentSummary(
                name="no-chunks",
                summary_fields=[
                    Summary(name="id"),
                    Summary(name="content"),
                    Summary(name="loc"),
                    Summary(name="title"),
                ],
            )
        )

        # Document summary: top_n_chunks (select top N chunks based on similarity)
        schema.add_document_summary(
            DocumentSummary(
                name="top_k_chunks",
                from_disk=True,
                summary_fields=[
                    Summary(name="id"),
                    Summary(name="loc"),
                    Summary(name="title"),
                    Summary(
                        name="chunks_topk",
                        fields=[
                            ("source", "chunks"),
                            (
                                "select-elements-by",
                                "top_k_chunk_combined_scores",
                            ),
                        ],
                    ),
                ],
            )
        )

    def add_rank_profile(self, schema: Schema) -> None:
        # Common inputs
        common_inputs = [
            ("query(embedding)", f"tensor<float>(x[{self.embedding_dim}])"),
            ("query(k)", "double"),
        ]

        # Functions for layered ranking: coarse content relevance, then best chunks
        chunk_functions = [
            Function(
                name="chunk_text_scores",
                expression=("map(bm25(chunks), f(x)((2 / 3.141592653589793) * atan(x)))"),
            ),
            Function(
                name="chunk_emb_vecs",
                expression="attribute(chunk_embeddings)",
            ),
            Function(name="content_vec", expression="attribute(content_embedding)"),
            Function(
                name="vector_norms",
                expression="sqrt(sum(pow(t, 2), x))",
                args=["t"],
            ),
            Function(
                name="chunk_sim_scores",
                expression=(
                    "reduce(query(embedding) * chunk_emb_vecs(), sum, x) / "
                    "(vector_norms(chunk_emb_vecs()) * vector_norms(query(embedding)))"
                ),
            ),
            Function(
                name="content_sim_score",
                expression=(
                    "reduce(query(embedding) * content_vec(), sum, x) / "
                    "(vector_norms(content_vec()) * vector_norms(query(embedding)))"
                ),
            ),
            Function(
                name="content_text_score",
                expression="(2 / 3.141592653589793) * atan(bm25(content))",
            ),
            Function(
                name="chunk_combined_scores",
                expression=("0.5 * chunk_sim_scores() + " "0.3 * chunk_text_scores() + " "0.2 * content_sim_score()"),
            ),
            Function(
                name="chunk_count",
                expression="if(query(k) > 0, query(k), 3)",
            ),
            Function(
                name="top_k_chunk_combined_scores",
                expression="top(chunk_count(), chunk_combined_scores())",
            ),
            Function(
                name="best_chunk_score",
                expression="reduce(top_k_chunk_combined_scores(), max, chunk)",
            ),
        ]

        summary_features = [
            "top_k_chunk_combined_scores",
            "best_chunk_score",
            "content_sim_score",
            "content_text_score",
        ]

        schema.add_rank_profile(
            RankProfile(
                name="default",
                inputs=common_inputs,
                functions=chunk_functions,
                summary_features=summary_features,
                rank_properties=[("query(chunks).elementGap", "1")],
                first_phase="sum(chunk_combined_scores())",
                second_phase=SecondPhaseRanking(expression="best_chunk_score()", rerank_count=100),
            )
        )

    def get_package(self) -> ApplicationPackage:
        """Get the Vespa application package with schema and rank profile."""
        schema = self.create_schema_fields()
        self.add_rank_profile(schema)
        app_package = ApplicationPackage(name=self.app_package_name, schema=[schema])
        return app_package

    def save_package(self, output_dir: str) -> None:
        """Save the Vespa application package to the specified directory."""
        app_package = self.get_package()
        app_package.to_files(output_dir)
