import os
import chromadb
from chromadb.config import Settings

# Engineering knowledge base regarding semiconductor wafer defects
KNOWLEDGE_BASE = [
    {
        "id": "defect_edge_ring",
        "title": "Edge-Ring Defect Troubleshooting",
        "content": "Edge-Ring defects typically appear as a continuous ring of failing dies around the outer edge of the wafer. Common Root Causes: 1. Uneven gas distribution in the etching chamber. 2. Non-uniform chuck temperature during deposition or etching. 3. Edge-bead removal issues during photolithography. Recommended Action: Inspect gas flow regulators and recalibrate chuck temperature sensors. Schedule maintenance for edge-bead removal module."
    },
    {
        "id": "defect_center",
        "title": "Center Defect Troubleshooting",
        "content": "Center defects are concentrated in the middle of the wafer. Common Root Causes: 1. Poor spin-coating uniformity (photoresist pooling in the center). 2. Center-heavy deposition profile. 3. Excessive center heating on the electrostatic chuck. Recommended Action: Verify spin speed and acceleration in the coating track. Check gas showerhead for clogging in the center region."
    },
    {
        "id": "defect_scratch",
        "title": "Scratch Defect Troubleshooting",
        "content": "Scratch defects manifest as linear patterns of failing dies, often crossing the wafer. Common Root Causes: 1. Mechanical handling damage by robotic arms or end-effectors. 2. Particulate contamination causing dragging during CMP (Chemical Mechanical Polishing). 3. Cassette or FOUP abrasion. Recommended Action: Check robot alignment and end-effector cleanliness. Inspect CMP pad conditioning and slurry filtration system."
    },
    {
        "id": "defect_donut",
        "title": "Donut Defect Troubleshooting",
        "content": "Donut defects appear as a ring, but not at the very edge (like Edge-Ring), leaving the center and extreme edge relatively clean. Common Root Causes: 1. Radially dependent temperature non-uniformity during rapid thermal processing (RTP). 2. Specific gas flow dynamics creating standing waves or depletion zones in the chamber. Recommended Action: Recalibrate RTP lamp zones. Inspect gas showerhead and exhaust pumping symmetry."
    },
    {
        "id": "general_forecast_strategy",
        "title": "Material Forecast & Yield Strategy",
        "content": "When the predicted material waste percentage rises above 5%, the factory must proactively increase raw material orders (wafers, photoresist, precursor gases) for the next quarter to compensate for the lower yield. High fail rates typically necessitate a temporary slow-down of production throughput to allow for deep tool maintenance and recalibration."
    }
]

def ingest_data():
    print("Initializing ChromaDB Persistent Client...")
    db_path = os.path.join(os.path.dirname(__file__), "chroma_db")
    client = chromadb.PersistentClient(path=db_path)
    
    # Create or get collection
    collection = client.get_or_create_collection(
        name="semiconductor_knowledge",
        metadata={"hnsw:space": "cosine"}
    )
    
    # Clear existing data if any (for idempotency)
    existing_ids = collection.get()['ids']
    if existing_ids:
        collection.delete(ids=existing_ids)
    
    # Prepare data for insertion
    ids = [item['id'] for item in KNOWLEDGE_BASE]
    documents = [item['content'] for item in KNOWLEDGE_BASE]
    metadatas = [{"title": item['title']} for item in KNOWLEDGE_BASE]
    
    print(f"Adding {len(documents)} documents to the knowledge base...")
    collection.add(
        documents=documents,
        metadatas=metadatas,
        ids=ids
    )
    
    print("Ingestion complete. ChromaDB is ready.")

if __name__ == "__main__":
    ingest_data()