#!/usr/bin/env python3
"""Generate the intentionally flawed Parquet fixtures used by the LLM audit article.

The fixtures are tiny on purpose. They are teaching data, not representative training
data, and every seeded problem is recorded in the adjacent JSON manifest.

For a downloaded copy, keep ``llm-training-audit.sql`` and
``rag-embedding-audit.sql`` beside this script, then run::

    python -m pip install duckdb==1.5.3
    python generate_llm_audit_fixtures.py

The fixtures and manifest are written beside the downloaded script. In this
repository, the source script writes to ``public/data/seo`` instead.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import duckdb


SCRIPT_PATH = Path(__file__).resolve()
TRAINING_SQL_BUNDLE_NAME = "llm-training-audit.sql"
RAG_SQL_BUNDLE_NAME = "rag-embedding-audit.sql"


def resolve_output_dir() -> Path:
    """Use the published artifact directory in-repo and the script directory elsewhere."""

    script_dir = SCRIPT_PATH.parent
    repository_root = script_dir.parent
    if script_dir.name == "scripts" and (repository_root / "package.json").is_file():
        return repository_root / "public" / "data" / "seo"
    return script_dir


OUTPUT_DIR = resolve_output_dir()
GENERATED_ON = "2026-08-24"

TRAINING_FIXTURE = OUTPUT_DIR / "llm-training-audit-fixture.parquet"
RAG_FIXTURE = OUTPUT_DIR / "rag-embedding-audit-fixture.parquet"
MANIFEST = OUTPUT_DIR / "llm-parquet-audit-manifest.json"
PUBLIC_GENERATOR = OUTPUT_DIR / "generate_llm_audit_fixtures.py"
TRAINING_SQL_BUNDLE = OUTPUT_DIR / TRAINING_SQL_BUNDLE_NAME
RAG_SQL_BUNDLE = OUTPUT_DIR / RAG_SQL_BUNDLE_NAME


TRAINING_AUDIT_SQL = r"""
SELECT
  count(*) AS rows,
  count_if(example_id IS NULL OR trim(example_id) = '') AS missing_example_ids,
  count(*) FILTER (WHERE example_id IS NOT NULL AND trim(example_id) <> '')
    - count(DISTINCT example_id) FILTER (WHERE example_id IS NOT NULL AND trim(example_id) <> '')
    AS duplicate_id_excess_rows,
  count_if(text IS NULL OR trim(text) = '') AS blank_text_rows,
  count_if(split IS NULL OR split NOT IN ('train', 'validation', 'test')) AS invalid_split_rows,
  count_if(source_uri IS NULL OR trim(source_uri) = '') AS missing_source_rows,
  count_if(license IS NULL OR trim(license) = '') AS missing_license_rows,
  count_if(token_ids IS NULL OR array_length(token_ids) = 0) AS empty_token_rows,
  count_if(tokenizer_revision IS NULL OR tokenizer_revision <> 'demo-tokenizer-v1')
    AS unexpected_tokenizer_rows
FROM read_parquet(?);
"""

TRAINING_LEAKAGE_SQL = r"""
WITH normalized AS (
  SELECT
    split,
    sha256(regexp_replace(lower(trim(text)), '\s+', ' ', 'g')) AS text_digest
  FROM read_parquet(?)
  WHERE text IS NOT NULL AND trim(text) <> ''
), cross_split AS (
  SELECT text_digest
  FROM normalized
  GROUP BY text_digest
  HAVING count(DISTINCT split) > 1
)
SELECT count(*) AS cross_split_text_groups
FROM cross_split;
"""

RAG_AUDIT_SQL = r"""
WITH numbered AS (
  SELECT row_number() OVER () AS audit_row_number, *
  FROM read_parquet(?)
), vector_stats AS (
  SELECT
    audit_row_number,
    count_if(value IS NULL OR NOT isfinite(value)) AS non_finite_values,
    sum(value * value) AS squared_norm
  FROM numbered, UNNEST(embedding) AS vector(value)
  GROUP BY audit_row_number
)
SELECT
  count(*) AS rows,
  count_if(d.chunk_id IS NULL OR trim(d.chunk_id) = '') AS missing_chunk_ids,
  count(*) FILTER (WHERE d.chunk_id IS NOT NULL AND trim(d.chunk_id) <> '')
    - count(DISTINCT d.chunk_id)
      FILTER (WHERE d.chunk_id IS NOT NULL AND trim(d.chunk_id) <> '')
    AS duplicate_chunk_id_excess_rows,
  count_if(d.embedding IS NULL) AS missing_embedding_rows,
  count_if(d.embedding IS NOT NULL AND array_length(d.embedding) <> 3) AS wrong_dimension_rows,
  count_if(coalesce(v.non_finite_values, 0) > 0) AS non_finite_rows,
  count_if(
    array_length(d.embedding) = 3
    AND coalesce(v.non_finite_values, 0) = 0
    AND v.squared_norm = 0
  ) AS zero_vector_rows,
  count_if(
    d.embedding_model_revision IS NULL
    OR d.embedding_model_revision <> 'demo-embed-v1'
  ) AS unexpected_model_rows
FROM numbered AS d
LEFT JOIN vector_stats AS v USING (audit_row_number);
"""


def sql_path(path: Path) -> str:
    return path.as_posix().replace("'", "''")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def result_dict(connection: duckdb.DuckDBPyConnection, sql: str, params: list[str]) -> dict[str, int]:
    cursor = connection.execute(sql, params)
    row = cursor.fetchone()
    if row is None:
        raise RuntimeError("Audit query returned no result")
    return {description[0]: int(value) for description, value in zip(cursor.description, row)}


def schema(connection: duckdb.DuckDBPyConnection, path: Path) -> list[dict[str, str]]:
    rows = connection.execute(
        "DESCRIBE SELECT * FROM read_parquet(?)",
        [str(path)],
    ).fetchall()
    return [
        {"name": row[0], "type": row[1], "nullable": row[2]}
        for row in rows
    ]


def main() -> None:
    OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
    missing_bundles = [
        bundle.name
        for bundle in (TRAINING_SQL_BUNDLE, RAG_SQL_BUNDLE)
        if not bundle.exists()
    ]
    if missing_bundles:
        raise FileNotFoundError(
            f"Missing companion SQL bundle(s): {', '.join(missing_bundles)}. "
            "Keep them beside the downloaded generator script."
        )

    connection = duckdb.connect()

    connection.execute(
        """
        CREATE OR REPLACE TABLE training_fixture AS
        SELECT *
        FROM (VALUES
          ('train-001', 'doc-a', 'train',
           'A Parquet row group stores column chunks for a subset of rows.',
           'https://example.test/docs/parquet-row-groups', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 24, 67, 9]),
          ('train-002', 'doc-b', 'train',
           'Whitespace normalization is useful, but it does not detect paraphrases.',
           'https://example.test/docs/normalization', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 31, 88, 9]),
          ('train-003', 'doc-c', 'validation',
           '  a parquet ROW group stores column chunks for a subset of rows.  ',
           'https://example.test/docs/copied-row-groups', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 24, 67, 9]),
          ('train-004', 'doc-d', 'validation', '   ',
           'https://example.test/docs/blank', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 9]),
          ('train-005', 'doc-e', 'archive',
           'A split value outside the contract should block release.',
           'https://example.test/docs/splits', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 45, 73, 9]),
          ('train-006', 'doc-f', 'train',
           'Provenance and license fields are operational data, not decoration.',
           NULL, NULL, 'demo-tokenizer-v1', [101, 12, 90, 9]),
          ('train-002', 'doc-g', 'train',
           'Duplicate example identifiers make reconciliation ambiguous.',
           'https://example.test/docs/identifiers', 'CC-BY-4.0',
           'demo-tokenizer-v1', [101, 63, 9]),
          ('train-008', 'doc-h', 'test',
           'Token arrays must be tied to one exact tokenizer revision.',
           'https://example.test/docs/tokens', 'CC-BY-4.0',
           'demo-tokenizer-v2', [])
        ) AS rows(
          example_id, source_document_id, split, text, source_uri, license,
          tokenizer_revision, token_ids
        );
        """
    )
    connection.execute(
        f"COPY training_fixture TO '{sql_path(TRAINING_FIXTURE)}' "
        "(FORMAT PARQUET, COMPRESSION ZSTD);"
    )

    connection.execute(
        """
        CREATE OR REPLACE TABLE rag_fixture AS
        SELECT *
        FROM (VALUES
          ('chunk-001', 'doc-a', 0, 'A row group is a horizontal partition of rows.',
           'text-v1', 'https://example.test/docs/row-groups', 'CC-BY-4.0',
           'demo-embed-v1', [0.10::DOUBLE, 0.20, 0.30]),
          ('chunk-002', 'doc-b', 0, 'A zero vector is numerically finite but often unusable.',
           'text-v1', 'https://example.test/docs/vectors', 'CC-BY-4.0',
           'demo-embed-v1', [0.0::DOUBLE, 0.0, 0.0]),
          ('chunk-003', 'doc-c', 0, 'This vector has the wrong dimensionality.',
           'text-v1', 'https://example.test/docs/dimensions', 'CC-BY-4.0',
           'demo-embed-v1', [0.30::DOUBLE, 0.40]),
          ('chunk-004', 'doc-d', 1, 'Equal dimensions do not prove equal embedding models.',
           'text-v1', 'https://example.test/docs/model-revisions', 'CC-BY-4.0',
           'demo-embed-v2', [0.40::DOUBLE, 0.50, 0.60]),
          ('chunk-005', 'doc-e', 0, 'Non-finite values can poison similarity calculations.',
           'text-v1', 'https://example.test/docs/non-finite', 'CC-BY-4.0',
           'demo-embed-v1', [CAST('NaN' AS DOUBLE), 0.20, 0.30]),
          ('chunk-006', 'doc-f', 0, 'Chunk identifiers should join back to an exact text revision.',
           'text-v1', 'https://example.test/docs/lineage', 'CC-BY-4.0',
           'demo-embed-v1', [0.60::DOUBLE, 0.70, 0.80])
        ) AS rows(
          chunk_id, source_document_id, chunk_index, text, text_revision,
          source_uri, license, embedding_model_revision, embedding
        );
        """
    )
    connection.execute(
        f"COPY rag_fixture TO '{sql_path(RAG_FIXTURE)}' "
        "(FORMAT PARQUET, COMPRESSION ZSTD);"
    )

    training_findings = result_dict(connection, TRAINING_AUDIT_SQL, [str(TRAINING_FIXTURE)])
    training_findings.update(
        result_dict(connection, TRAINING_LEAKAGE_SQL, [str(TRAINING_FIXTURE)])
    )
    rag_findings = result_dict(
        connection,
        RAG_AUDIT_SQL,
        [str(RAG_FIXTURE)],
    )

    PUBLIC_GENERATOR.write_text(SCRIPT_PATH.read_text(encoding="utf-8"), encoding="utf-8")

    evidence_artifacts = [
        {
            "name": PUBLIC_GENERATOR.name,
            "role": "fixture_generator",
            "bytes": PUBLIC_GENERATOR.stat().st_size,
            "sha256": sha256(PUBLIC_GENERATOR),
        },
        {
            "name": TRAINING_SQL_BUNDLE.name,
            "role": "training_sql_bundle",
            "bytes": TRAINING_SQL_BUNDLE.stat().st_size,
            "sha256": sha256(TRAINING_SQL_BUNDLE),
        },
        {
            "name": RAG_SQL_BUNDLE.name,
            "role": "rag_sql_bundle",
            "bytes": RAG_SQL_BUNDLE.stat().st_size,
            "sha256": sha256(RAG_SQL_BUNDLE),
        },
    ]

    manifest = {
        "purpose": "Intentionally flawed teaching fixtures for the viewparquet LLM data audit article",
        "generated_on": GENERATED_ON,
        "generator": "scripts/generate_llm_audit_fixtures.py",
        "duckdb_version": connection.execute("SELECT version()").fetchone()[0],
        "evidence_artifacts": evidence_artifacts,
        "limitations": [
            "Synthetic teaching data; not representative of a production corpus",
            "Example tokenizer IDs and embedding values do not belong to a real model",
            "Passing these checks does not establish consent, legal fitness, semantic quality, or model suitability",
        ],
        "files": [
            {
                "name": TRAINING_FIXTURE.name,
                "bytes": TRAINING_FIXTURE.stat().st_size,
                "sha256": sha256(TRAINING_FIXTURE),
                "schema": schema(connection, TRAINING_FIXTURE),
                "expected_findings": training_findings,
            },
            {
                "name": RAG_FIXTURE.name,
                "bytes": RAG_FIXTURE.stat().st_size,
                "sha256": sha256(RAG_FIXTURE),
                "schema": schema(connection, RAG_FIXTURE),
                "expected_findings": rag_findings,
            },
        ],
    }
    MANIFEST.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
    print(json.dumps(manifest, indent=2))


if __name__ == "__main__":
    main()