from __future__ import annotations

import argparse
import json
from typing import Any

from embed_1c_semantic_cache import DEFAULT_ADAPTER_URL, adapter_call, batched, embedding_model_label
from rag_embedding_providers import LOCAL_HASHING_MODEL, LOCAL_HASHING_PROVIDER, embed_texts, provider_metadata


def code_embedding_model_label(*, provider: str, model: str, dimensions: int) -> str:
    label = embedding_model_label(provider=provider, model=model)
    normalized_provider = str(provider or "").strip().lower().replace("_", "-")
    if normalized_provider in {"openai-compatible", "openai"} and int(dimensions or 0) > 0:
        return f"{label}@d{int(dimensions)}"
    return label


def embed_pending_code_vectors(
    *,
    adapter_url: str,
    base_id: str,
    limit: int = 100,
    batch_size: int = 16,
    embedding_provider: str = LOCAL_HASHING_PROVIDER,
    embedding_model: str = LOCAL_HASHING_MODEL,
    dimensions: int = 384,
    embedding_base_url: str = "",
    embedding_api_key_env: str = "OPENAI_API_KEY",
    chunk_kinds: tuple[str, ...] | list[str] = ("routine",),
    max_text_chars: int = 4000,
    dry_run: bool = False,
    timeout_seconds: int = 180,
) -> dict[str, Any]:
    stored_model = code_embedding_model_label(
        provider=embedding_provider,
        model=embedding_model,
        dimensions=dimensions,
    )
    pending = adapter_call(
        adapter_url,
        "metadata.code_vector.pending",
        {
            "base_id": base_id,
            "embedding_model": stored_model,
            "limit": int(limit or 100),
            "chunk_kinds": list(dict.fromkeys(str(value).strip().lower() for value in chunk_kinds if str(value).strip())),
            "max_text_chars": int(max_text_chars),
        },
        timeout_seconds=timeout_seconds,
    )
    if pending.get("status") != "ok":
        return {
            "schema": "onec_code_vector_embedding_worker.v1",
            "status": pending.get("status") or "error",
            "error": pending.get("error"),
            "pending": pending,
        }
    chunks = [item for item in pending.get("chunks") or [] if isinstance(item, dict)]
    upserts: list[dict[str, Any]] = []
    skipped: list[dict[str, Any]] = []
    for batch in batched(chunks, max(int(batch_size or 1), 1)):
        texts = [str(item.get("text") or "") for item in batch]
        vectors = embed_texts(
            texts,
            provider=embedding_provider,
            model=embedding_model,
            dimensions=dimensions,
            base_url=embedding_base_url,
            api_key_env=embedding_api_key_env,
            timeout_seconds=timeout_seconds,
        )
        for item, vector in zip(batch, vectors):
            chunk_id = str(item.get("chunk_id") or "")
            text_sha1 = str(item.get("text_sha1") or "")
            if not chunk_id or not text_sha1 or not vector:
                skipped.append(
                    {
                        "chunk_id": chunk_id or None,
                        "reason": "missing_chunk_id_text_sha1_or_embedding",
                    }
                )
                continue
            if dry_run:
                upserts.append(
                    {
                        "status": "dry_run",
                        "chunk_id": chunk_id,
                        "text_sha1": text_sha1,
                        "dimensions": len(vector),
                    }
                )
                continue
            result = adapter_call(
                adapter_url,
                "metadata.code_vector.embedding.upsert",
                {
                    "base_id": base_id,
                    "chunk_id": chunk_id,
                    "text_sha1": text_sha1,
                    "embedding_model": stored_model,
                    "embedding": vector,
                },
                timeout_seconds=timeout_seconds,
            )
            upserts.append(
                {
                    "status": result.get("status"),
                    "error": result.get("error"),
                    "chunk_id": chunk_id,
                    "text_sha1": text_sha1,
                    "dimensions": result.get("dimensions") or len(vector),
                }
            )
    observed_dimensions = next(
        (
            int(item.get("dimensions") or 0)
            for item in upserts
            if int(item.get("dimensions") or 0) > 0
        ),
        int(dimensions),
    )
    provider = provider_metadata(
        provider=embedding_provider,
        model=embedding_model,
        dimensions=observed_dimensions,
        base_url=embedding_base_url,
    )
    return {
        "schema": "onec_code_vector_embedding_worker.v1",
        "status": "ok",
        "base_id": base_id,
        "adapter_url": adapter_url,
        "dry_run": bool(dry_run),
        "embedding": {
            "provider": provider.get("embedding_provider"),
            "model": embedding_model,
            "stored_embedding_model": stored_model,
            "dimensions": observed_dimensions,
            "chunk_kinds": list(chunk_kinds),
            "max_text_chars": int(max_text_chars),
        },
        "counts": {
            "pending": len(chunks),
            "processed": len(upserts),
            "stored": len([item for item in upserts if item.get("status") == "ok"]),
            "conflicts": len([item for item in upserts if item.get("status") == "conflict"]),
            "skipped": len(skipped),
            "errors": len(
                [
                    item
                    for item in upserts
                    if item.get("status") not in {"ok", "dry_run", "conflict"}
                ]
            ),
        },
        "upserts": upserts,
        **({"skipped": skipped} if skipped else {}),
    }


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Embed pending BSL code chunks from the 1C adapter local code index."
    )
    parser.add_argument("--adapter-url", default=DEFAULT_ADAPTER_URL)
    parser.add_argument("--base-id", required=True)
    parser.add_argument("--limit", type=int, default=100)
    parser.add_argument("--batch-size", type=int, default=16)
    parser.add_argument(
        "--embedding-provider",
        default=LOCAL_HASHING_PROVIDER,
        choices=[LOCAL_HASHING_PROVIDER, "openai-compatible"],
    )
    parser.add_argument("--embedding-model", default=LOCAL_HASHING_MODEL)
    parser.add_argument("--dimensions", type=int, default=384)
    parser.add_argument("--embedding-base-url", default="")
    parser.add_argument("--embedding-api-key-env", default="OPENAI_API_KEY")
    parser.add_argument(
        "--chunk-kind",
        action="append",
        choices=["routine", "module"],
        default=None,
        help="Chunk kind to embed; repeat to include both. Defaults to routine.",
    )
    parser.add_argument(
        "--max-text-chars",
        type=int,
        default=4000,
        help="Skip oversized chunks in this pass. Defaults to 4000 characters.",
    )
    parser.add_argument("--timeout-seconds", type=int, default=180)
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--json", action="store_true")
    args = parser.parse_args()

    result = embed_pending_code_vectors(
        adapter_url=args.adapter_url,
        base_id=args.base_id,
        limit=args.limit,
        batch_size=args.batch_size,
        embedding_provider=args.embedding_provider,
        embedding_model=args.embedding_model,
        dimensions=args.dimensions,
        embedding_base_url=args.embedding_base_url,
        embedding_api_key_env=args.embedding_api_key_env,
        chunk_kinds=tuple(args.chunk_kind or ["routine"]),
        max_text_chars=args.max_text_chars,
        dry_run=args.dry_run,
        timeout_seconds=args.timeout_seconds,
    )
    if args.json:
        print(json.dumps(result, ensure_ascii=False, indent=2))
    else:
        counts = result.get("counts") or {}
        print(
            "code vector embeddings: "
            f"pending={counts.get('pending')} processed={counts.get('processed')} "
            f"stored={counts.get('stored')} conflicts={counts.get('conflicts')} "
            f"errors={counts.get('errors')}"
        )
    return (
        0
        if result.get("status") == "ok"
        and int((result.get("counts") or {}).get("errors") or 0) == 0
        else 1
    )


if __name__ == "__main__":
    raise SystemExit(main())
