# app.py

import os
import uuid
import json
import streamlit as st
from dotenv import load_dotenv
from azure.search.documents.indexes import SearchIndexClient
from azure.search.documents.indexes.models import (
    SearchIndex, SimpleField, SearchableField, SearchField,
    SearchFieldDataType, VectorSearch, VectorSearchProfile,
    HnswAlgorithmConfiguration, SemanticConfiguration, SemanticPrioritizedFields,
    SemanticField, SemanticSearch
)
from azure.search.documents.models import VectorizedQuery
from azure.search.documents import SearchClient
from azure.core.credentials import AzureKeyCredential
from openai import AzureOpenAI

# --------------------- ENV ---------------------
load_dotenv()

AZURE_SEARCH_SERVICE_ENDPOINT = os.getenv("AZURE_SEARCH_SERVICE_ENDPOINT")
AZURE_SEARCH_INDEX = os.getenv("AZURE_SEARCH_INDEX")
AZURE_SEARCH_ADMIN_KEY = os.getenv("AZURE_SEARCH_ADMIN_KEY")

AZURE_OPENAI_ENDPOINT = os.getenv("AZURE_OPENAI_ENDPOINT")
AZURE_OPENAI_KEY = os.getenv("AZURE_OPENAI_KEY")
AZURE_OPENAI_EMBEDDING_DEPLOYMENT = os.getenv("AZURE_OPENAI_EMBEDDING_DEPLOYMENT")
AZURE_OPENAI_API_VERSION = os.getenv("AZURE_OPENAI_API_VERSION")
AZURE_OPENAI_EMBEDDING_DIMENSIONS = int(os.getenv("AZURE_OPENAI_EMBEDDING_DIMENSIONS", "1536"))

# --------------------- CLIENTS ---------------------
credential = AzureKeyCredential(AZURE_SEARCH_ADMIN_KEY)

search_client = SearchClient(
    endpoint=AZURE_SEARCH_SERVICE_ENDPOINT,
    index_name=AZURE_SEARCH_INDEX,
    credential=credential
)

index_client = SearchIndexClient(
    endpoint=AZURE_SEARCH_SERVICE_ENDPOINT,
    credential=credential
)

openai_client = AzureOpenAI(
    azure_endpoint=AZURE_OPENAI_ENDPOINT,
    api_key=AZURE_OPENAI_KEY,
    azure_deployment=AZURE_OPENAI_EMBEDDING_DEPLOYMENT,
    api_version=AZURE_OPENAI_API_VERSION
)

# --------------------- FUNCTIONS ---------------------
def create_index():
    try:
        index_client.get_index(AZURE_SEARCH_INDEX)
        return  # Index already exists
    except:
        pass

    fields = [
        SimpleField(name="id", type=SearchFieldDataType.String, key=True),
        SimpleField(name="userID", type=SearchFieldDataType.Int64, filterable=True),
        SearchableField(name="userName", type=SearchFieldDataType.String, filterable=True),
        SearchableField(name="type", type=SearchFieldDataType.String),
        SearchableField(name="description", type=SearchFieldDataType.String),
        SearchableField(name="lineNumber", type=SearchFieldDataType.String),
        SearchableField(name="text", type=SearchFieldDataType.String),
        SearchableField(name="timestamp", type=SearchFieldDataType.String),
        SearchField(
            name="embedding",
            type=SearchFieldDataType.Collection(SearchFieldDataType.Single),
            searchable=True,
            vector_search_dimensions=AZURE_OPENAI_EMBEDDING_DIMENSIONS,
            vector_search_profile_name="myHnswProfile"
        ),
    ]

    vector_search = VectorSearch(
        algorithms=[HnswAlgorithmConfiguration(name="myHnsw")],
        profiles=[VectorSearchProfile(name="myHnswProfile", algorithm_configuration_name="myHnsw")]
    )

    semantic_config = SemanticConfiguration(
        name="my-semantic-config",
        prioritized_fields=SemanticPrioritizedFields(
            title_field=SemanticField(field_name="userName"),
            content_fields=[SemanticField(field_name="text")]
        )
    )

    index = SearchIndex(
        name=AZURE_SEARCH_INDEX,
        fields=fields,
        vector_search=vector_search,
        semantic_search=SemanticSearch(configurations=[semantic_config])
    )

    index_client.create_or_update_index(index)

def generate_embeddings(texts):
    response = openai_client.embeddings.create(input=texts, model=AZURE_OPENAI_EMBEDDING_DEPLOYMENT)
    return [e.embedding for e in response.data]

def upload_documents(data, tracker_path="user_upload_tracker.json"):
    # Load or create tracking file
    if os.path.exists(tracker_path):
        with open(tracker_path, "r") as f:
            tracker = json.load(f)
    else:
        tracker = []

    enriched = []
    upload_counter = {}  # For current upload session

    for item in data:
        user_id = str(item.get("userId", "unknown"))
        user_name = item.get("userName", "unknown")
        key = (user_id, user_name)

        # Prepare enriched document
        enriched_item = {
            "id": str(uuid.uuid4()),
            "userID": int(user_id),
            "userName": user_name,
            "type": item.get("type", ""),
            "description": item.get("description", ""),
            "lineNumber": item.get("lineNumber", ""),
            "text": item.get("text", ""),
            "timestamp": item.get("timestamp", ""),
            "embedding": generate_embeddings([item.get("text", "")])[0]
        }

        enriched.append(enriched_item)

        if key not in upload_counter:
            upload_counter[key] = 0
        upload_counter[key] += 1

    # Upload documents to Azure Search
    search_client.upload_documents(enriched)

    # Update tracker
    for (user_id, user_name), count in upload_counter.items():
        found = False
        for record in tracker:
            if record["userId"] == user_id and record["userName"] == user_name:
                record["indexes"] += count
                found = True
                break
        if not found:
            tracker.append({
                "userId": user_id,
                "userName": user_name,
                "indexes": count
            })

    # Save updated tracker
    with open(tracker_path, "w") as f:
        json.dump(tracker, f, indent=2)

    return tracker, upload_counter

def search_query(query_text, top_k=5):
    vector = generate_embeddings([query_text])[0]
    results = search_client.search(
        vector_queries=[
            VectorizedQuery(
                vector=vector,
                k=top_k,
                fields="embedding"
            )
        ]
    )
    return results

# --------------------- STREAMLIT UI ---------------------
st.set_page_config("Macy's Notes Search", layout="wide")
st.title("🔎 Macy's Notes Search — Azure Native")

uploaded_file = st.file_uploader("Upload Macy's JSON Notes", type="json")

if uploaded_file:
    st.info("Creating index if not exists...")
    create_index()

    data = json.load(uploaded_file)
    st.info("Generating embeddings and uploading data...")
    tracker, current_uploads = upload_documents(data)

    st.success("✅ Data uploaded and indexed successfully!")

    st.subheader("📦 Current Upload Summary (from this file):")
    for (user_id, user_name), count in current_uploads.items():
        st.markdown(f"- **{user_name}** (`{user_id}`) → `{count}` new records")

    st.subheader("📊 Cumulative Upload Summary (all time):")
    for entry in tracker:
        st.markdown(f"- **{entry['userName']}** (`{entry['userId']}`) → **Total records indexed:** `{entry['indexes']}`")

query = st.text_input("Ask a question about the notes:")

if query:
    st.write("🔍 Searching with vector embeddings...")
    results = search_query(query)
    for i, r in enumerate(results):
        st.markdown(f"**{i+1}.** {r['text']}  \n*Timestamp:* `{r['timestamp']}`")



AZURE_OPENAI_DEPLOYMENT = os.getenv("AZURE_OPENAI_DEPLOYMENT")

# ------------------ CLIENTS ------------------
credential = AzureKeyCredential(AZURE_SEARCH_ADMIN_KEY)

search_client = SearchClient(
    endpoint=AZURE_SEARCH_SERVICE_ENDPOINT,
    index_name=AZURE_SEARCH_INDEX,
    credential=credential
)

openai_client = AzureOpenAI(
    azure_endpoint=AZURE_OPENAI_ENDPOINT,
    api_key=AZURE_OPENAI_KEY,
    azure_deployment=AZURE_OPENAI_DEPLOYMENT,
    api_version=AZURE_OPENAI_API_VERSION
)

# ------------------ FUNCTIONS ------------------
def search_by_user_id(user_id, top_k):
    filter_str = f"userID eq {user_id}"
    results = search_client.search(
        search_text="*",
        filter=filter_str,
        top=top_k
    )
    print(f"Found {len(results)} documents for user ID {user_id}.", "Top K:", top_k)
    results = []
    for result in results:
        if "text" in result:
            results.append({
                "id": result["id"],
                "text": result["text"],
                "userID": result["userID"]
            })
    return list(results)

def summarize_texts(texts, model="gpt-4"):
    prompt = "Summarize the following notes according to the order id :\n\n" + "\n\n".join(texts)
    response = openai_client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": "You are a helpful summarizer."},
            {"role": "user", "content": prompt}
        ],
        temperature=0
    )
    return response.choices[0].message.content

# ------------------ STREAMLIT UI ------------------
st.set_page_config("🔍 Macy's Notes Retrieval", layout="wide")
st.title("📘 Macy's Notes Summarizer")

# Load tracker
tracker_path = "user_upload_tracker.json"

if os.path.exists(tracker_path):
    with open(tracker_path, "r") as f:
        tracker_data = json.load(f)

    user_ids = [entry["userId"] for entry in tracker_data]
    
    # Filter the list of user IDs based on user input (start typing and options will filter)
    search_input = st.text_input("Start typing User ID:", "")
    filtered_user_ids = [user_id for user_id in user_ids if search_input in str(user_id)]

    if filtered_user_ids:
        selected_user_id = st.selectbox("Select User ID", filtered_user_ids)
    else:
        selected_user_id = None

    if selected_user_id:
        selected_user = next((u for u in tracker_data if u["userId"] == selected_user_id), None)

        st.markdown(f"👤 **User Name:** `{selected_user['userName']}`")
        st.markdown(f"📌 **Total Indexed Records:** `{selected_user['indexes']}`")

        if st.button("🔍 Summarize Notes"):
            with st.spinner("Fetching and summarizing notes..."):
                user_docs = search_by_user_id(selected_user_id, selected_user["indexes"])
                user_texts = [doc["text"] for doc in user_docs if doc.get("text")]
                if not user_texts:
                    st.warning("No documents found for this user.")
                else:
                    summary = summarize_texts(user_texts)
                    st.subheader("🧾 Summary:")
                    st.markdown(summary)
else:
    st.error("❌ No tracking file (`user_upload_tracker.json`) found. Please run upload first.")


but what is the requirement is each time the user upload the file the previous data from the azure cognitive search must be removed and from our file where we record also must be removed and what we can filter with UserId and we got the n number of index from our files and inside it also user can give order id that we can search from the index got from UserId and then display the summary