# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import json
import re
from dataclasses import dataclass
from typing import Dict, Mapping, Optional, Sequence

import pyarrow as pa

from pypaimon.common.options.core_options import CoreOptions
from pypaimon.multimodal.query import (
    BatchVectorQuery,
    HybridQuery,
    ScanQuery,
    TextQuery,
    VectorQuery,
)
from pypaimon.schema.data_types import PyarrowFieldParser, is_blob_type
from pypaimon.table.data_evolution_merge_into import (
    WhenMatched,
    WhenNotMatched,
    source_col as _source_col,
)
from pypaimon.table.special_fields import SpecialFields


_ALL_SOURCE_COLUMNS = object()
_MERGE_REF_PATTERN = re.compile(r'\b(source|target)\.(\w+)\b')
_STRING_LITERAL_PATTERN = re.compile(r"'(?:[^']|'')*'")


@dataclass(frozen=True)
class VectorRoute:
    """Vector route spec for hybrid search."""

    column: Optional[str]
    vector: object
    weight: float = 1.0
    limit: Optional[int] = None
    options: Optional[Dict[str, str]] = None


@dataclass(frozen=True)
class TextRoute:
    """Full-text route spec for hybrid search."""

    query: object
    column: Optional[str] = None
    weight: float = 1.0
    limit: Optional[int] = None
    options: Optional[Dict[str, str]] = None


def vector_route(column, vector, *, weight: float = 1.0,
                 limit: Optional[int] = None,
                 options: Optional[Dict[str, str]] = None) -> VectorRoute:
    return VectorRoute(
        column=column,
        vector=vector,
        weight=weight,
        limit=limit,
        options=dict(options or {}),
    )


def text_route(query, *, column: Optional[str] = None, weight: float = 1.0,
               limit: Optional[int] = None,
               options: Optional[Dict[str, str]] = None) -> TextRoute:
    return TextRoute(
        query=query,
        column=column,
        weight=weight,
        limit=limit,
        options=dict(options or {}),
    )


class MultimodalTable:
    """High-level table facade for mutable multimodal application data."""

    def __init__(self, catalog, identifier: str, raw_table):
        self.catalog = catalog
        self.identifier = identifier
        self.name = identifier
        self.raw_table = raw_table
        self.table = raw_table

    def add(self, data):
        arrow_table = _to_arrow_table(data, _target_schema(self.raw_table))
        write_builder = self.raw_table.new_batch_write_builder()
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        try:
            table_write.write_arrow(arrow_table)
            table_commit.commit(table_write.prepare_commit())
        finally:
            table_write.close()
            table_commit.close()
        return self

    def overwrite(self, data, partition: Optional[Mapping[str, object]] = None):
        arrow_table = _to_arrow_table(data, _target_schema(self.raw_table))
        overwrite_partition = dict(partition) if partition is not None else None
        write_builder = (
            self.raw_table.new_batch_write_builder()
            .overwrite(overwrite_partition)
        )
        table_write = write_builder.new_write()
        table_commit = write_builder.new_commit()
        try:
            table_write.write_arrow(arrow_table)
            table_commit.commit(table_write.prepare_commit())
        finally:
            table_write.close()
            table_commit.close()
        return self

    def update(self, where, values):
        query = self.scan().where(where)
        predicate = query._predicate
        write_builder = self.raw_table.new_batch_write_builder()
        table_update = write_builder.new_update()
        table_commit = write_builder.new_commit()
        try:
            messages = table_update.update_by_predicate(predicate, values)
            table_commit.commit(messages)
        finally:
            table_commit.close()
        return self

    def delete(self, where):
        query = self.scan().where(where)
        predicate = query._predicate
        write_builder = self.raw_table.new_batch_write_builder()
        table_update = write_builder.new_update()
        table_commit = write_builder.new_commit()
        try:
            messages = table_update.delete_by_predicate(predicate)
            table_commit.commit(messages)
        finally:
            table_commit.close()
        return self

    def merge(self, on):
        return _MergeBuilder(self, on)

    def map_with_blobs(self, dataset, columns, fn, **kwargs):
        from pypaimon.ray import map_with_blobs

        return map_with_blobs(
            dataset,
            columns,
            fn,
            file_io=self.raw_table.file_io,
            all_blob_columns=_blob_columns(self.raw_table),
            **kwargs,
        )

    def scan(
            self,
            *,
            snapshot_id: Optional[int] = None,
            tag_name: Optional[str] = None):
        return ScanQuery(_time_travel_table(
            self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name))

    def take_row_ids(
            self,
            row_ids,
            *,
            snapshot_id: Optional[int] = None,
            tag_name: Optional[str] = None):
        row_ids = _coerce_row_ids(row_ids)
        read_table = _time_travel_table(
            self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name)
        read_builder = read_table.new_read_builder().with_projection(
            [field.name for field in read_table.fields]
            + [SpecialFields.ROW_ID.name]
        )
        predicate = read_builder.new_predicate_builder().is_in(
            SpecialFields.ROW_ID.name, row_ids)
        query = ScanQuery(read_table)
        query._predicate = predicate
        return query

    def blobs(self, *, column: Optional[str] = None, key_column: Optional[str] = None):
        from pypaimon.multimodal.blob_store import BlobStore
        return BlobStore(self, column=column, key_column=key_column)

    def search(
            self,
            query,
            *,
            column: Optional[str] = None,
            options: Optional[Dict[str, str]] = None,
            pre_filter=None,
            snapshot_id: Optional[int] = None,
            tag_name: Optional[str] = None):
        read_table = _time_travel_table(
            self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name)
        schema = _target_schema(read_table)
        if isinstance(query, str):
            return TextQuery(
                read_table,
                text_query=_coerce_full_text_query(
                    query, "search", schema, column=column),
                pre_filter=pre_filter,
            )
        vector = _coerce_vector(query, "search")
        if _looks_like_batch_vectors(vector):
            raise ValueError(
                "search() accepts a single query; use search_vectors() for "
                "multiple vectors.")
        return VectorQuery(
            read_table,
            vector=vector,
            vector_column=column or _infer_vector_column(schema, "column"),
            vector_options=options,
            pre_filter=pre_filter,
        )

    def search_vectors(
            self,
            vectors,
            *,
            column: Optional[str] = None,
            options: Optional[Dict[str, str]] = None,
            pre_filter=None,
            snapshot_id: Optional[int] = None,
            tag_name: Optional[str] = None):
        read_table = _time_travel_table(
            self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name)
        schema = _target_schema(read_table)
        vectors = _coerce_vectors(vectors)
        vector_column = column or _infer_vector_column(schema, "column")
        return BatchVectorQuery(
            read_table,
            vectors=vectors,
            vector_column=vector_column,
            vector_options=options,
            pre_filter=pre_filter,
        )

    def search_hybrid(
            self,
            routes,
            *,
            ranker: str = "rrf",
            route_limit: Optional[int] = None,
            pre_filter=None,
            snapshot_id: Optional[int] = None,
            tag_name: Optional[str] = None):
        read_table = _time_travel_table(
            self.raw_table, snapshot_id=snapshot_id, tag_name=tag_name)
        schema = _target_schema(read_table)
        vector_routes, text_routes = _normalize_hybrid_routes(
            schema,
            routes=routes,
            method="search_hybrid",
        )
        if not vector_routes and not text_routes:
            raise ValueError(
                "search_hybrid requires at least one route.")
        return HybridQuery(
            read_table,
            vector_routes=vector_routes,
            text_routes=text_routes,
            ranker=ranker,
            route_limit=route_limit,
            pre_filter=pre_filter,
        )

    def create_index(self, column, index_type, options=None):
        return self.raw_table.create_global_index(
            column, index_type=_normalize_index_type(index_type), options=options)

    def _merge(
            self,
            source,
            on,
            when_matched: Sequence[WhenMatched],
            when_not_matched: Sequence[WhenNotMatched],
            operation: str):
        target_cols, source_cols = _normalize_merge_on(on, operation)
        target_schema = _target_schema(self.raw_table)
        source_table = _to_arrow_table(source)
        _validate_merge_on(
            target_schema, source_table, target_cols, source_cols, operation)
        source_table = _cast_merge_on_columns(
            source_table, target_schema, target_cols, source_cols)
        when_matched, when_not_matched = _resolve_merge_clauses(
            when_matched,
            when_not_matched,
            target_schema,
            source_table,
            dict(zip(target_cols, source_cols)),
        )

        write_builder = self.raw_table.new_batch_write_builder()
        table_update = write_builder.new_update()
        table_commit = write_builder.new_commit()
        try:
            messages = table_update.merge_into(
                source_table,
                on=dict(zip(target_cols, source_cols)),
                when_matched=list(when_matched),
                when_not_matched=list(when_not_matched),
            )
            table_commit.commit(messages)
        finally:
            table_commit.close()
        return self


class _MergeBuilder:
    """Builder for idempotent merge writes."""

    def __init__(self, table: MultimodalTable, on):
        self._table = table
        self._on = on
        self._when_matched = []
        self._when_not_matched = []

    def when_matched_update(self, values=None, where: Optional[str] = None):
        self._when_matched.append(
            WhenMatched.update(
                _ALL_SOURCE_COLUMNS if values is None else values,
                condition=_normalize_merge_where(where),
            ))
        return self

    def when_matched_delete(self, where: Optional[str] = None):
        self._when_matched.append(
            WhenMatched.delete(
                condition=_normalize_merge_where(where),
            ))
        return self

    def when_not_matched_insert(self, values=None, where: Optional[str] = None):
        self._when_not_matched.append(
            WhenNotMatched(
                insert=_ALL_SOURCE_COLUMNS if values is None else values,
                condition=_normalize_merge_where(where),
            ))
        return self

    def execute(self, data):
        if not self._when_matched and not self._when_not_matched:
            raise ValueError(
                "merge requires at least one matched or not-matched clause.")
        return self._table._merge(
            data,
            self._on,
            when_matched=self._when_matched,
            when_not_matched=self._when_not_matched,
            operation="merge",
        )


def _blob_columns(table):
    return tuple(
        field.name for field in table.fields
        if is_blob_type(field.type)
    )


def _to_arrow_table(data, target_schema=None):
    if isinstance(data, pa.Table):
        table = data
    elif isinstance(data, pa.RecordBatch):
        table = pa.Table.from_batches([data])
    elif isinstance(data, list):
        table = pa.Table.from_pylist(data)
    elif isinstance(data, dict):
        table = pa.Table.from_pydict(data)
    elif hasattr(data, "__dataframe__") or data.__class__.__module__.startswith("pandas"):
        table = pa.Table.from_pandas(data, preserve_index=False)
    else:
        raise ValueError("Unsupported multimodal data type: %r" % type(data))

    if target_schema is None:
        return table
    return _align_to_schema(table, target_schema)


def _coerce_row_ids(row_ids):
    if row_ids is None or isinstance(row_ids, (str, bytes)):
        raise ValueError("row_ids must be an iterable of row id integers.")
    try:
        iterator = iter(row_ids)
    except TypeError:
        raise ValueError("row_ids must be an iterable of row id integers.")

    coerced = []
    for row_id in iterator:
        if hasattr(row_id, "as_py"):
            row_id = row_id.as_py()
        coerced.append(int(row_id))
    return coerced


def _time_travel_table(table, snapshot_id=None, tag_name=None):
    if snapshot_id is not None and tag_name is not None:
        raise ValueError(
            "snapshot_id and tag_name cannot be set at the same time")
    if snapshot_id is not None:
        return table.copy({
            CoreOptions.SCAN_SNAPSHOT_ID.key(): str(snapshot_id),
        })
    if tag_name is not None:
        return table.copy({
            CoreOptions.SCAN_TAG_NAME.key(): tag_name,
        })
    return table


def _target_schema(table):
    return PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)


def _align_to_schema(
        table: pa.Table,
        schema: pa.Schema,
        column_mapping: Optional[Dict[str, str]] = None) -> pa.Table:
    column_mapping = column_mapping or {}
    arrays = []
    for field in schema:
        source_name = column_mapping.get(field.name, field.name)
        if source_name in table.column_names:
            array = table[source_name]
            if array.type != field.type:
                array = array.cast(field.type)
        else:
            array = pa.nulls(table.num_rows, type=field.type)
        arrays.append(array)
    return pa.Table.from_arrays(arrays, schema=schema)


def _normalize_merge_on(on, operation: str):
    if isinstance(on, Mapping):
        target_cols = list(on.keys())
        source_cols = list(on.values())
    elif isinstance(on, str):
        target_cols = [on]
        source_cols = [on]
    else:
        target_cols = list(on)
        source_cols = list(on)
    if not target_cols:
        raise ValueError("%s requires at least one ON column." % operation)
    if len(target_cols) != len(source_cols):
        raise ValueError(
            "%s ON target/source column counts do not match." % operation)
    return target_cols, source_cols


def _validate_merge_on(
        target_schema,
        source_table,
        target_cols,
        source_cols,
        operation: str):
    target_names = set(target_schema.names)
    source_names = set(source_table.column_names)
    missing_targets = [col for col in target_cols if col not in target_names]
    if missing_targets:
        raise ValueError(
            "%s ON target columns are not in table schema: %s"
            % (operation, missing_targets))
    missing_sources = [col for col in source_cols if col not in source_names]
    if missing_sources:
        raise ValueError(
            "%s ON source columns are not in source data: %s"
            % (operation, missing_sources))


def _cast_merge_on_columns(source_table, target_schema, target_cols, source_cols):
    result = source_table
    for target_col, source_col_name in zip(target_cols, source_cols):
        target_type = target_schema.field(target_col).type
        index = result.schema.get_field_index(source_col_name)
        if index < 0:
            continue
        array = result[source_col_name]
        if array.type == target_type:
            continue
        source_field = result.schema.field(index)
        result = result.set_column(
            index,
            pa.field(
                source_col_name,
                target_type,
                nullable=source_field.nullable,
                metadata=source_field.metadata,
            ),
            array.cast(target_type),
        )
    return result


def _resolve_merge_clauses(
        when_matched,
        when_not_matched,
        target_schema,
        source_table,
        on_map):
    target_names = list(target_schema.names)
    source_names = set(source_table.column_names)
    return (
        [
            (
                WhenMatched.delete(condition=clause.condition)
                if clause.delete
                else WhenMatched.update(
                    _resolve_merge_set_spec(
                        clause.update, target_names, source_names, on_map),
                    condition=clause.condition,
                )
            )
            for clause in when_matched
        ],
        [
            WhenNotMatched(
                insert=_resolve_merge_set_spec(
                    clause.insert, target_names, source_names, on_map),
                condition=clause.condition,
            )
            for clause in when_not_matched
        ],
    )


def _resolve_merge_set_spec(spec, target_names, source_names, on_map):
    if spec is not _ALL_SOURCE_COLUMNS:
        return spec
    return {
        target_name: _source_col(source_name)
        for target_name in target_names
        for source_name in [on_map.get(target_name, target_name)]
        if source_name in source_names
    }


def _normalize_merge_where(where):
    if where is None:
        return None
    if not isinstance(where, str):
        raise ValueError("merge where must be a string.")
    parts, last = [], 0
    for match in _STRING_LITERAL_PATTERN.finditer(where):
        parts.append(_rewrite_merge_refs(where[last:match.start()]))
        parts.append(match.group())
        last = match.end()
    parts.append(_rewrite_merge_refs(where[last:]))
    return ''.join(parts)


def _rewrite_merge_refs(text):
    return _MERGE_REF_PATTERN.sub(
        lambda m: "%s.%s" % (
            "s" if m.group(1) == "source" else "t",
            m.group(2),
        ),
        text,
    )


def _normalize_index_type(index_type):
    if not isinstance(index_type, str):
        return index_type
    normalized = index_type.strip().lower()
    if normalized == "full-text":
        return "full-text"
    return index_type


def _coerce_vector(value, method):
    if value is None or isinstance(value, str):
        raise ValueError("%s requires a vector query." % method)
    if hasattr(value, "tolist"):
        value = value.tolist()
    try:
        return list(value)
    except TypeError as e:
        raise ValueError("%s requires a vector query." % method) from e


def _coerce_vectors(values):
    method = "search_vectors"
    if values is None or isinstance(values, str):
        raise ValueError("%s requires query vectors." % method)
    if hasattr(values, "tolist"):
        values = values.tolist()
    try:
        vectors = list(values)
    except TypeError as e:
        raise ValueError("%s requires query vectors." % method) from e
    if not vectors:
        raise ValueError("%s requires at least one vector query." % method)
    return [_coerce_vector(vector, method) for vector in vectors]


def _looks_like_batch_vectors(value):
    if not value:
        return False
    first = value[0]
    return not isinstance(first, (int, float)) and hasattr(first, "__iter__")


def _normalize_hybrid_routes(
        schema,
        routes,
        method):
    vector_routes, text_routes = [], []
    for route in _route_specs(routes):
        if _is_text_route_spec(route):
            text_routes.append(_normalize_text_route(route, method, schema))
        else:
            vector_routes.append(
                _normalize_vector_route(route, method, schema))
    return vector_routes, text_routes


def _route_specs(routes):
    if routes is None:
        return []
    if isinstance(routes, (VectorRoute, TextRoute, Mapping)):
        return [routes]
    if _is_column_vector_pair(routes):
        return [routes]
    return list(routes)


def _is_text_route_spec(route):
    if isinstance(route, TextRoute):
        return True
    if isinstance(route, VectorRoute):
        return False
    if isinstance(route, Mapping):
        route_type = route.get("type") or route.get("route_type")
        if route_type is not None:
            return str(route_type).replace("-", "_").lower() in (
                "text", "full_text", "fulltext")
        return (
            "query" in route
            or "text" in route
            or "terms" in route
            or "full_text_query" in route
        ) and not _has_vector_query(route)
    return False


def _normalize_vector_route(route, method, schema):
    if isinstance(route, VectorRoute):
        column = route.column
        vector = route.vector
        weight = route.weight
        limit = route.limit
        options = route.options
    elif isinstance(route, Mapping):
        column = (
            route.get("column")
            or route.get("vector_column")
            or route.get("field")
            or route.get("anns_field")
        )
        vector = _mapping_vector(route)
        weight = route.get("weight", 1.0)
        limit = route.get("limit")
        options = _mapping_options(route)
    elif _is_column_vector_pair(route):
        column, vector = route
        weight = 1.0
        limit = None
        options = None
    else:
        raise ValueError(
            "%s vector routes require a route spec with a column and vector."
            % method)
    if not column:
        raise ValueError("%s vector routes require a column." % method)
    return {
        "column": column,
        "vector": _coerce_vector(vector, method),
        "weight": weight,
        "limit": limit,
        "options": dict(options or {}),
    }


def _normalize_text_route(route, method, schema):
    if isinstance(route, TextRoute):
        column = route.column
        query = route.query
        weight = route.weight
        limit = route.limit
        options = route.options
    elif isinstance(route, Mapping):
        column = route.get("column")
        query = (
            route.get("query")
            or route.get("text")
            or route.get("terms")
            or route.get("full_text_query")
        )
        weight = route.get("weight", 1.0)
        limit = route.get("limit")
        options = route.get("options")
    else:
        raise ValueError(
            "%s text routes require a route spec with a query."
            % method)
    text_query = _coerce_full_text_query(query, method, schema, column=column)
    return {
        "column": text_query["column"],
        "query": text_query["query"],
        "weight": weight,
        "limit": limit,
        "options": dict(options or {}),
    }


def _has_vector_query(route):
    return any(key in route for key in ("vector", "query_vector", "data"))


def _mapping_vector(route):
    if "vector" in route:
        return route["vector"]
    if "query_vector" in route:
        return route["query_vector"]
    if "data" in route:
        data = route["data"]
        if hasattr(data, "tolist"):
            data = data.tolist()
        values = list(data)
        if len(values) == 1 and hasattr(values[0], "__iter__"):
            return values[0]
        return values
    raise ValueError("Vector route requires vector, query_vector, or data.")


def _mapping_options(route):
    if "options" in route:
        return route["options"]
    if "param" in route:
        param = route["param"]
        if isinstance(param, Mapping) and isinstance(param.get("params"), Mapping):
            return param["params"]
        return param
    if "params" in route:
        return route["params"]
    return None


def _is_column_vector_pair(route):
    return (
        isinstance(route, (list, tuple))
        and len(route) == 2
        and isinstance(route[0], str)
    )


def _coerce_full_text_query(query, method, schema, column=None):
    field_name = column or _infer_text_column(schema, "text")
    if isinstance(query, str):
        stripped = query.lstrip()
        query_json = (
            query
            if stripped.startswith("{")
            else json.dumps({"match": {"query": query}}, separators=(",", ":"))
        )
        return {"column": field_name, "query": query_json}
    if isinstance(query, Mapping):
        return {
            "column": field_name,
            "query": json.dumps(query, separators=(",", ":")),
        }
    raise ValueError("%s requires a text string or query mapping." % method)


def _infer_vector_column(schema: pa.Schema, parameter: str = "vector_column"):
    columns = [
        field.name
        for field in schema
        if pa.types.is_fixed_size_list(field.type)
    ]
    return _infer_single_column(columns, "vector", parameter)


def _infer_text_column(schema: pa.Schema, parameter: str = "text_column"):
    columns = [
        field.name
        for field in schema
        if pa.types.is_string(field.type) or pa.types.is_large_string(field.type)
    ]
    return _infer_single_column(columns, "text", parameter)


def _infer_single_column(columns, kind, parameter):
    if len(columns) == 1:
        return columns[0]
    if not columns:
        raise ValueError("No %s column found; pass %s." % (kind, parameter))
    raise ValueError(
        "Multiple %s columns found %s; pass %s."
        % (kind, columns, parameter))
