"""UPSERT query builder - sync version."""
import numpy as np
from typing import Dict, List, Union, Any, Optional, Tuple
from veclite.query.ir import UpsertIR
from veclite.query.sqlgen import generate_upsert
from veclite.query.dialect.sqlite import SQLiteDialect
from veclite.core.results import Result


class SyncUpsertBuilder:
    """Fluent builder for UPSERT (INSERT ... ON CONFLICT) queries - sync version."""

    def __init__(
        self,
        client,
        schema,
        table: str,
        values: Union[Dict[str, Any], List[Dict[str, Any]]],
        on_conflict: Union[str, List[str]],
        ignore_duplicates: bool = False,
        returning: str = "representation",
        count: Optional[str] = None,
        default_to_null: bool = False,
    ):
        self.db = client
        self.schema = schema
        self.table = table
        self.dialect = SQLiteDialect()
        self.rows = [values] if isinstance(values, dict) else values
        self.do_nothing = ignore_duplicates
        self.returning_all = returning == "representation"

        if self.table in schema.views:
            raise ValueError(f"Cannot UPSERT into view '{table}'.")

        if isinstance(on_conflict, str):
            self.on_conflict_cols = [c.strip() for c in on_conflict.split(",") if c.strip()]
        else:
            self.on_conflict_cols = list(on_conflict)

        tbl = self.schema.get_table(table)
        fields = tbl.get_fields()
        for c in self.on_conflict_cols:
            if c not in fields:
                raise ValueError(f"on_conflict column '{c}' does not exist in table '{table}'")

        if len(self.on_conflict_cols) == 1:
            fd = fields[self.on_conflict_cols[0]]
            if not (fd.primary_key or fd.unique):
                raise ValueError(
                    f"on_conflict column '{self.on_conflict_cols[0]}' must have UNIQUE constraint or be PRIMARY KEY."
                )
        else:
            declared_uniques = [tuple(u) for u in getattr(tbl, '__uniques__', ())]
            if tuple(self.on_conflict_cols) not in declared_uniques:
                raise ValueError(
                    f"on_conflict {tuple(self.on_conflict_cols)} does not match any declared composite UNIQUE."
                )

    def _validate_upsert_row(self, row: Dict[str, Any]) -> Dict[str, Any]:
        """Validate upsert data with smart handling for updates vs inserts

        For primary key upserts (e.g., on_conflict="id"), we:
        - Still validate types and field-level constraints for provided fields
        - Skip the "missing required fields" check since PK upserts are always updates

        For other upserts, we do full validation including required fields.
        """
        if not isinstance(row, dict):
            raise ValueError(f"Upsert data for table '{self.table}' must be a dictionary")

        table_cls = self.schema.get_table(self.table)
        table_instance = table_cls()
        table_instance._schema = self.schema

        table_fields = table_instance.get_fields()

        is_pk_update = any(
            table_fields.get(col, None) and table_fields[col].primary_key
            for col in self.on_conflict_cols
        )

        validated_data = {}

        for field_name, value in row.items():
            if field_name in table_fields:
                field_desc = table_fields[field_name]

                if field_desc.primary_key and field_name not in self.on_conflict_cols:
                    raise ValueError(
                        f"Cannot manually set primary key field '{field_name}' in table '{self.table}'. "
                        f"Primary keys are auto-generated."
                    )

                field_desc.set_context(self.table, field_name, self.schema)
                validated_data[field_name] = field_desc.validate(value)
            else:
                raise ValueError(f"Field '{field_name}' not found in table '{self.table}'")

        if not is_pk_update:
            table_instance._validate_required_fields(validated_data, table_fields)

        return validated_data

    def _build_pk_upsert_with_select(
        self,
        row: Dict[str, Any],
        provided_cols: List[str],
        missing_required_cols: set,
        pk_col: str
    ) -> Tuple[str, List[Any]]:
        """Build SELECT-based INSERT for PK upserts with missing required fields.

        When upserting by primary key with partial data, we need to pull
        missing NOT NULL fields from the existing row to avoid constraint violations.
        """
        all_cols = provided_cols + list(missing_required_cols)
        select_parts = []
        params = []

        for col in all_cols:
            if col in provided_cols:
                select_parts.append("?")
                params.append(row[col])
            else:
                select_parts.append(self.dialect.q(col))

        params.append(row[pk_col])

        update_clauses = [
            f"{self.dialect.q(c)} = excluded.{self.dialect.q(c)}"
            for c in provided_cols if c != pk_col
        ]
        action = f"DO UPDATE SET {', '.join(update_clauses)}" if update_clauses else "DO NOTHING"

        sql = (
            f"INSERT INTO {self.dialect.q(self.table)} "
            f"({', '.join(self.dialect.q(c) for c in all_cols)}) "
            f"SELECT {', '.join(select_parts)} "
            f"FROM {self.dialect.q(self.table)} "
            f"WHERE {self.dialect.q(pk_col)} = ? "
            f"ON CONFLICT ({self.dialect.q(pk_col)}) {action}"
        )

        if self.returning_all:
            sql += " RETURNING *"

        return sql, params

    def execute(self):
        validated_rows = [self._validate_upsert_row(row) for row in self.rows]
        with_defaults = [self.db._apply_runtime_defaults(self.table, row) for row in validated_rows]
        serialized_rows = [self.db._serialize_json_fields(self.table, row) for row in with_defaults]

        if not serialized_rows:
            return Result([])

        table_cls = self.schema.get_table(self.table)
        table_fields = table_cls.get_fields() if table_cls else {}

        is_pk_conflict = any(
            table_fields.get(col, None) and table_fields[col].primary_key
            for col in self.on_conflict_cols
        )

        required_fields = set()
        for field_name, field_desc in table_fields.items():
            if not field_desc.nullable and field_desc.default is None and not field_desc.primary_key:
                required_fields.add(field_name)

        provided_fields = set(serialized_rows[0].keys())
        missing_required = required_fields - provided_fields

        use_select_insert = is_pk_conflict and missing_required and self.returning_all

        cols = list(serialized_rows[0].keys())
        num_cols = len(cols)
        max_vars = self.db._get_max_vars()

        from veclite.query.validation import check_vector_embedder_required
        has_vector_fields = check_vector_embedder_required(
            self.table, self.schema, self.db.embedder, "UPSERT into"
        )

        total_params = num_cols * len(serialized_rows)
        use_executemany = total_params > max_vars and not has_vector_fields and not use_select_insert

        def _exec_upsert():
            all_results = []

            if use_select_insert:
                for row in serialized_rows:
                    sql, params = self._build_pk_upsert_with_select(
                        row, cols, missing_required, self.on_conflict_cols[0]
                    )
                    rows = self.db._exec(sql, params)
                    for r in rows:
                        r = self.db._deserialize_json_fields(self.table, r)
                        all_results.append(r)

            elif use_executemany:
                conflict_cols = ", ".join(self.dialect.q(c) for c in self.on_conflict_cols)
                if self.do_nothing:
                    action = "DO NOTHING"
                else:
                    update_clauses = [
                        f"{self.dialect.q(c)} = excluded.{self.dialect.q(c)}"
                        for c in cols if c not in self.on_conflict_cols
                    ]
                    action = f"DO UPDATE SET {', '.join(update_clauses)}" if update_clauses else "DO NOTHING"

                sql = (
                    f"INSERT INTO {self.dialect.q(self.table)} ({', '.join(self.dialect.q(c) for c in cols)}) "
                    f"VALUES ({', '.join(['?'] * num_cols)}) "
                    f"ON CONFLICT ({conflict_cols}) {action}"
                )
                param_rows = [[row[c] for c in cols] for row in serialized_rows]
                with self.db._lock:
                    self.db.conn.executemany(sql, param_rows)
                    self.db.conn.commit()
                return None
            else:
                batch_size = max(1, max_vars // num_cols)
                for i in range(0, len(serialized_rows), batch_size):
                    batch = serialized_rows[i:i + batch_size]
                    ir = UpsertIR(
                        table=self.table,
                        rows=batch,
                        on_conflict=self.on_conflict_cols,
                        do_nothing=self.do_nothing,
                        returning_all=self.returning_all,
                    )
                    sql, params = generate_upsert(ir, self.dialect)
                    rows = self.db._exec(sql, params)

                    for row in rows:
                        row = self.db._deserialize_json_fields(self.table, row)
                        all_results.append(row)
            return all_results

        all_results = _exec_upsert()

        if all_results is None:
            return Result([], count=len(serialized_rows))

        self._embed_vectors(all_results)

        return Result(all_results)

    def _embed_vectors(self, rows: List[Dict[str, Any]]):
        if not rows or not self.db.embedder:
            return

        table_cls = self.schema.get_table(self.table)
        if not table_cls:
            return

        input_fields = set()
        for row in self.rows:
            input_fields.update(row.keys())

        vector_fields = []
        for field_name, field in table_cls.get_fields().items():
            if getattr(field, 'vector', False) and field_name in input_fields:
                vector_fields.append(field_name)

        if not vector_fields:
            return

        for field_name in vector_fields:
            field = table_cls.get_fields()[field_name]
            is_contextualized = getattr(field, 'contextualized', False)

            texts = []
            ids = []
            for row in rows:
                if field_name in row and row[field_name]:
                    texts.append(row[field_name])
                    ids.append(row['id'])

            if not texts:
                continue

            # Use appropriate embedding method
            if is_contextualized:
                # Contextualized: each text is a separate document
                # Pass as [[text1], [text2], ...] to preserve document boundaries
                inputs = [[text] for text in texts]
                nested_embeddings = self.db.embedder.contextualized_embed(
                    inputs=inputs,
                    model="voyage-context-3",
                    input_type="document",
                    output_dimension=getattr(field, 'contextualized_dim', 512)
                )
                # Flatten nested results: [[emb1], [emb2]] -> [emb1, emb2]
                embeddings = [doc_embs[0] for doc_embs in nested_embeddings]
            else:
                # Regular vector field
                embeddings = self.db.embedder.embed(texts)

            vectors = [np.array(emb, dtype=np.float32) for emb in embeddings]

            vector_store = self.db.get_or_create_vector_store(self.table, field_name)
            vector_store.add_batch(ids, vectors)
