"""
JAIM - Data Augmentation & Re-indexing Script
Generates 4 paraphrased symptom variants per entry using Groq (Llama 3.3 70B),
embeds everything with BAAI/bge-large-en-v1.5 (1024 dims),
and batch-upserts all vectors into Pinecone.

Run this ONCE before starting the app.
Requires: A Pinecone index with dimension=1024, metric=cosine.
"""

import os
import json
import time
import requests
from dotenv import load_dotenv
from pinecone import Pinecone
from sentence_transformers import SentenceTransformer
from embed_pdf import extract_entries_from_pdf

load_dotenv()

# ─── Config ──────────────────────────────────────────────────────────────────
PINECONE_API_KEY = os.getenv("PINECONE_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")
INDEX_NAME = os.getenv("PINECONE_INDEX_NAME", "jaim3")
EMBEDDING_MODEL = "BAAI/bge-large-en-v1.5"
GROQ_MODEL = "llama-3.3-70b-versatile"
PDF_PATH = "db.pdf"
NUM_PARAPHRASES = 4
BATCH_SIZE = 50


# ─── Paraphrase Generation via Groq ──────────────────────────────────────────
def generate_paraphrases(symptoms: str) -> list[str]:
    """Use Groq (Llama 3.3 70B) to generate paraphrased symptom variants."""
    url = "https://api.groq.com/openai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json",
    }

    prompt = f"""Generate exactly {NUM_PARAPHRASES} different paraphrased versions of the following plant symptom description.
Each paraphrase should:
- Use different vocabulary and sentence structure
- Preserve the exact medical/botanical meaning
- Sound like how a farmer or gardener would naturally describe the problem

Return ONLY a valid JSON array of {NUM_PARAPHRASES} strings. No extra text, no markdown.

Original symptoms: "{symptoms}"
"""

    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "system", "content": "You paraphrase text. Return only valid JSON arrays."},
            {"role": "user", "content": prompt},
        ],
        "temperature": 0.8,
        "max_tokens": 512,
    }

    for attempt in range(3):
        resp = requests.post(url, json=payload, headers=headers, timeout=30)
        if resp.status_code == 429:
            wait = 2 ** (attempt + 1)
            print(f"   ⏳ Rate limited, waiting {wait}s...")
            time.sleep(wait)
            continue
        resp.raise_for_status()
        break
    else:
        raise RuntimeError("Groq rate limit exceeded after 3 retries")

    content = resp.json()["choices"][0]["message"]["content"].strip()
    # Handle potential markdown code-block wrapping
    if content.startswith("```"):
        content = content.split("\n", 1)[1].rsplit("```", 1)[0].strip()

    paraphrases = json.loads(content)
    return paraphrases[:NUM_PARAPHRASES]


# ─── Main Pipeline ────────────────────────────────────────────────────────────
def main():
    print("🔄 Loading BAAI/bge-large-en-v1.5 embedding model (1024 dims)...")
    model = SentenceTransformer(EMBEDDING_MODEL)

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

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

    all_vectors = []

    for i, entry in enumerate(entries):
        original_id = f"entry_{i}_chunk_0"
        symptoms = entry["symptoms"]

        print(f"\n{'═' * 60}")
        print(f"📝 Entry {i} | Cause: {entry['cause_given']}")
        print(f"   Original: {symptoms[:80]}...")

        # ── Generate paraphrases via Groq ─────────────────────────
        print(f"   🤖 Generating {NUM_PARAPHRASES} paraphrases...")
        try:
            paraphrases = generate_paraphrases(symptoms)
        except Exception as e:
            print(f"   ⚠️ Failed to generate paraphrases: {e}")
            paraphrases = []
        time.sleep(1)  # Rate-limit buffer between Groq calls

        # ── Shared metadata (same for original + all variants) ────
        base_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"],
        }

        # ── Embed original symptoms ──────────────────────────────
        # BGE spec: documents do NOT need a prefix
        original_embedding = model.encode(symptoms).tolist()

        all_vectors.append({
            "id": original_id,
            "values": original_embedding,
            "metadata": {
                **base_metadata,
                "is_augmented": False,
                "parent_id": original_id,
            }
        })

        # ── Embed paraphrased variants ───────────────────────────
        for j, para in enumerate(paraphrases):
            aug_embedding = model.encode(para).tolist()
            all_vectors.append({
                "id": f"entry_{i}_aug_{j}",
                "values": aug_embedding,
                "metadata": {
                    **base_metadata,
                    "symptoms": para,       # Store paraphrased version
                    "is_augmented": True,
                    "parent_id": original_id,
                }
            })
            print(f"   ✅ Variant {j + 1}: {para[:80]}...")

    # ── Batch upsert to Pinecone ──────────────────────────────────
    total = len(all_vectors)
    print(f"\n📤 Uploading {total} vectors to Pinecone in batches of {BATCH_SIZE}...")
    for start in range(0, total, BATCH_SIZE):
        batch = all_vectors[start : start + BATCH_SIZE]
        index.upsert(vectors=batch)
        print(f"   Batch {start // BATCH_SIZE + 1}: {len(batch)} vectors uploaded")

    print("\n✅ Augmentation & re-indexing complete!")
    stats = index.describe_index_stats()
    print(f"📊 Index stats: {stats}")
    print(f"   Expected: {len(entries)} originals + {len(entries) * NUM_PARAPHRASES} augmented = {total} total")


if __name__ == "__main__":
    main()