from feast import Entity, FeatureView, Field, FileSource, PushSource
from feast.types import Array, Float32, String, Int64
from feast.value_type import ValueType
from datetime import timedelta

document = Entity(
    name="document_id",
    value_type=ValueType.STRING,
    description="Unique identifier for document chunks"
)

document_embeddings_source = FileSource(
    name="document_embeddings_source",
    path="data/document_embeddings.parquet",
    timestamp_field="event_timestamp",
    created_timestamp_column="created_timestamp",
)

document_embeddings_push_source = PushSource(
    name="document_embeddings_push_source",
    batch_source=document_embeddings_source,
)

document_embeddings = FeatureView(
    name="document_embeddings",
    entities=[document],
    ttl=timedelta(days=365),
    schema=[
        Field(name="embedding", dtype=Array(Float32), vector_index=True),
        Field(name="chunk_text", dtype=String),
        Field(name="document_title", dtype=String),
        Field(name="chunk_index", dtype=Int64),
        Field(name="file_path", dtype=String),
        Field(name="chunk_length", dtype=Int64),
    ],
    online=True,
    source=document_embeddings_push_source,
    tags={"team": "rag", "version": "v3"},
)

import pandas as pd
import numpy as np
from datetime import datetime
from sentence_transformers import SentenceTransformer
from feast import FeatureStore
from feast.data_format import PushMode

# Generate test embeddings (384 dimensions)
model = SentenceTransformer('all-MiniLM-L6-v2')
texts = [
    'Test document 1',
    'Test document 2', 
    'Test document 3',
    'Test document 4',
    'Test document 5'
]
embeddings = model.encode(texts)  # Shape: (5, 384)

# Create DataFrame
feature_data = []
for i, (text, embedding) in enumerate(zip(texts, embeddings)):
    feature_data.append({
        "document_id": f"test_doc_{i}",
        "embedding": embedding.tolist(),  # Convert to list as per docs
        "chunk_text": text,
        "document_title": "test_document.md",
        "chunk_index": i,
        "file_path": "test_path",
        "chunk_length": len(text),
        "event_timestamp": pd.Timestamp.now(tz='UTC'),
        "created_timestamp": pd.Timestamp.now(tz='UTC')
    })

df = pd.DataFrame(feature_data)
print(f"Input data: {len(df)} rows, {len(df) * 384} total elements")

# Initialize Feast store
fs = FeatureStore(repo_path="feast_feature_repo")

# This will fail with dimension mismatch
fs.push(
    push_source_name="document_embeddings_push_source",
    df=df,
    to=PushMode.ONLINE_AND_OFFLINE
)

# Save to parquet file
df.to_parquet('feast_feature_repo/data/document_embeddings.parquet', index=False)

# Try materialization
from datetime import timedelta
end_time = datetime.now()
start_time = end_time - timedelta(hours=1)

# This will also fail with same dimension mismatch
fs.materialize(
    start_date=start_time,
    end_date=end_time,
    feature_views=["document_embeddings"]
)