"""
JAIM - Embedding Pipeline
Parses the Vrikshayurveda PDF table and uploads embeddings to Pinecone.
This script is used to populate the Pinecone index with the treatment data.
Run this only if you need to re-embed the data.
"""

import os
import re
from dotenv import load_dotenv
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer
import PyPDF2

load_dotenv()

# ─── Config ──────────────────────────────────────────────────────────────────
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "JaimRAG")
EMBEDDING_MODEL = "all-mpnet-base-v2"
PDF_PATH = "db.pdf"

# ─── Parse the PDF table ─────────────────────────────────────────────────────
def extract_entries_from_pdf(pdf_path: str) -> list[dict]:
    """
    Extracts structured entries from the Vrikshayurveda PDF table.
    Each entry has: disorder, cause_given, symptoms, cause_elaborated,
    possible_causes, treatment_material.
    """
    reader = PyPDF2.PdfReader(pdf_path)
    full_text = ""
    for page in reader.pages:
        text = page.extract_text()
        if text:
            full_text += text + "\n"

    # Define the known disorder entries from the table
    entries = [
        {
            "disorder": "Internal",
            "cause_given": "Vata",
            "symptoms": "Trunk bent; knots on the trunk or leaves, hard fruits (less juicy and pleasant), slow defoliation, the loss of flowers and fruits, and general yellowing of the leaves and fruits.",
            "cause_elaborated": "Arid land",
            "possible_causes": ["Root-infecting fungi or nematodes", "viruses", "and saline/alkaline soils."],
            "treatment_material": "Application of fermented mixture of hog fat, porpoise oil, ghee (clarified butter), hemp, horsehair, and cow horn-boiled and set to decoction; also use of panchmula."
        },
        {
            "disorder": "",
            "cause_given": "Kapha",
            "symptoms": "Fruit bearing is delayed, and the fruits are bland and overripe, oozing without injuries.",
            "cause_elaborated": "Appear in winter and spring",
            "possible_causes": ["Fungal gummosis or rot", "nutrient deficiencies or toxicities", "and excessive watering."],
            "treatment_material": "Application of white mustard paste at the roots, followed by the watering of the trees with a sesame-and-ash mixture; the earth at the roots of the trees should then be removed and replaced with fresh, dry earth."
        },
        {
            "disorder": "",
            "cause_given": "Pitta",
            "symptoms": "Early leaf withering, or early fruit or flower decay.",
            "cause_elaborated": "Occur at the end of the summer.",
            "possible_causes": ["Viral diseases", "salinity in irrigation water", "and susceptibility to blossom blight and fruit decay caused by fungal or bacterial infections."],
            "treatment_material": "Sprinkling of cold water on the trees and application of a paste made of sesame, then sprinkled with milk and water. Water-stress is caused by the combination of milk water and crab shell smoke."
        },
        {
            "disorder": "External",
            "cause_given": "Struck by axe etc",
            "symptoms": "Trees wounded resulting in drying up.",
            "cause_elaborated": "NA",
            "possible_causes": ["Same as causes given."],
            "treatment_material": "Tree wounds can be healed by applying a paste made from the bark of Nyagrodha (Banyan tree) and Udumbara (Fig tree), along with cow dung, honey, and ghee."
        },
        {
            "disorder": "",
            "cause_given": "Faulty Seed",
            "symptoms": "Trees become unproductive.",
            "cause_elaborated": "Lack of appropriate seed treatment; wrong remedies used",
            "possible_causes": ["Seed infected with pathogens or infected by insects."],
            "treatment_material": "Seed should be treated with milk, mustard, ash of sesame and brhati, rubbing with cow dung, honey and/or bidanga."
        },
        {
            "disorder": "",
            "cause_given": "Ants",
            "symptoms": "Foul smell, original fragrance missing; reduction of leaf size, stunted seedlings.",
            "cause_elaborated": "NA",
            "possible_causes": ["Ants could mean a wide range of insects."],
            "treatment_material": "Worms (caterpillars) on trees can be removed by smoking with a mixture of white mustard, ramatha, vidanga, vaca, usana and water mixed with beef, pigeon flesh, billatta powder, and horn of a buffalo. Trees can also be anointed with vidanga mixed with ghee, watered for seven days with soft water, and treated with an ointment made of beef."
        },
        {
            "disorder": "",
            "cause_given": "Excessive watering",
            "symptoms": "Trees suffer from indigestion. Destruction of trees.",
            "cause_elaborated": "NA",
            "possible_causes": ["NA"],
            "treatment_material": "A mixture of honey and vidanga should be applied to every root of tender plants after being uprooted and scratched with nails."
        }
    ]

    return entries


# ─── Embed and Upload ─────────────────────────────────────────────────────────
def embed_and_upload():
    """Embed entries and upload to Pinecone."""
    print("🔄 Loading embedding model...")
    model = SentenceTransformer(EMBEDDING_MODEL)

    print("📄 Extracting entries from PDF...")
    entries = extract_entries_from_pdf(PDF_PATH)
    print(f"   Found {len(entries)} entries")

    # Connect to Pinecone
    print("🌲 Connecting to Pinecone...")
    pc = Pinecone(api_key=PINECONE_API_KEY)
    index = pc.Index(INDEX_NAME)

    # Create embeddings and upsert
    vectors = []
    for i, entry in enumerate(entries):
        # Create a rich text representation for embedding (Symptoms ONLY)
        text_for_embedding = f"Symptoms: {entry['symptoms']}"

        embedding = model.encode(text_for_embedding).tolist()

        vectors.append({
            "id": f"entry_{i}_chunk_0",
            "values": embedding,
            "metadata": {
                "disorder": entry["disorder"],
                "cause_given": entry["cause_given"],
                "symptoms": entry["symptoms"],
                "cause_elaborated": entry["cause_elaborated"],
                "possible_causes": entry["possible_causes"],
                "treatment_material": entry["treatment_material"]
            }
        })

    print(f"📤 Uploading {len(vectors)} vectors to Pinecone...")
    index.upsert(vectors=vectors)
    print("✅ Done! All entries embedded and uploaded.")

    # Verify
    stats = index.describe_index_stats()
    print(f"📊 Index stats: {stats}")


if __name__ == "__main__":
    embed_and_upload()