"""
RAG App: Amazon Bedrock + Aurora PostgreSQL (pgvector)
Answers questions grounded in your own Oracle migration docs.

Setup:
  1. Create an Aurora PostgreSQL cluster (see setup.sql)
  2. Set environment variables (see below)
  3. pip install -r requirements.txt
  4. python rag_app.py
"""

import boto3
import psycopg2
import json
import os
from typing import List

# ── Config ─────────────────────────────────────────────────────────────────────
DB_HOST     = os.environ.get("DB_HOST", "your-cluster.cluster-xxxx.us-east-1.rds.amazonaws.com")
DB_NAME     = os.environ.get("DB_NAME", "ragdb")
DB_USER     = os.environ.get("DB_USER", "postgres")
DB_PASSWORD = os.environ.get("DB_PASSWORD", "your-password")
AWS_REGION  = os.environ.get("AWS_REGION", "us-east-1")

EMBED_MODEL    = "amazon.titan-embed-text-v2:0"
GENERATE_MODEL = "anthropic.claude-sonnet-4-6"   # Claude Sonnet 4.6 via Bedrock

bedrock = boto3.client("bedrock-runtime", region_name=AWS_REGION)


# ── Database helpers ────────────────────────────────────────────────────────────
def get_db():
    return psycopg2.connect(
        host=DB_HOST, dbname=DB_NAME, user=DB_USER, password=DB_PASSWORD,
        connect_timeout=10
    )


def setup_db():
    """Enable pgvector and create the docs table."""
    with get_db() as conn:
        with conn.cursor() as cur:
            cur.execute("CREATE EXTENSION IF NOT EXISTS vector;")
            cur.execute("""
                CREATE TABLE IF NOT EXISTS docs (
                    id        SERIAL PRIMARY KEY,
                    content   TEXT    NOT NULL,
                    embedding vector(1024)
                );
            """)
            # ivfflat index for fast approximate nearest-neighbour search
            cur.execute("""
                CREATE INDEX IF NOT EXISTS docs_emb_idx
                ON docs USING ivfflat (embedding vector_cosine_ops)
                WITH (lists = 50);
            """)
        conn.commit()
    print("Aurora pgvector ready.")


# ── Bedrock helpers ─────────────────────────────────────────────────────────────
def embed(text: str) -> List[float]:
    """Convert text to a 1024-dim vector using Titan Embeddings v2."""
    resp = bedrock.invoke_model(
        modelId=EMBED_MODEL,
        body=json.dumps({"inputText": text, "dimensions": 1024}),
    )
    return json.loads(resp["body"].read())["embedding"]


# ── RAG functions ───────────────────────────────────────────────────────────────
def ingest(documents: List[str]):
    """Embed each document and store it in Aurora."""
    with get_db() as conn:
        with conn.cursor() as cur:
            for doc in documents:
                vec = embed(doc)
                cur.execute(
                    "INSERT INTO docs (content, embedding) VALUES (%s, %s::vector)",
                    (doc, vec)
                )
        conn.commit()
    print(f"Ingested {len(documents)} documents into Aurora pgvector.")


def retrieve(question: str, top_k: int = 3) -> List[str]:
    """Find the top-k most similar docs using cosine similarity."""
    q_vec = embed(question)
    with get_db() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT content
                FROM   docs
                ORDER  BY embedding <=> %s::vector
                LIMIT  %s
                """,
                (q_vec, top_k)
            )
            return [row[0] for row in cur.fetchall()]


def ask(question: str) -> str:
    """Full RAG loop: retrieve context → generate grounded answer with Claude."""
    print(f"Searching Aurora pgvector for relevant docs...")
    chunks = retrieve(question)

    context = "\n\n---\n\n".join(chunks)

    prompt = f"""You are a helpful assistant for Oracle DBAs migrating to Amazon Aurora PostgreSQL.
Use ONLY the context below to answer the question.
If the answer is not in the context, say "I don't have enough information."

Context:
{context}

Question: {question}

Answer:"""

    print("Sending context to Claude on Bedrock...")
    resp = bedrock.invoke_model(
        modelId=GENERATE_MODEL,
        body=json.dumps({
            "anthropic_version": "bedrock-2023-05-31",
            "max_tokens": 512,
            "messages": [{"role": "user", "content": prompt}]
        })
    )
    return json.loads(resp["body"].read())["content"][0]["text"]


# ── Sample docs ─────────────────────────────────────────────────────────────────
ORACLE_MIGRATION_DOCS = [
    """Oracle SEQUENCE vs Aurora PostgreSQL SEQUENCE
Oracle: standalone object, call with seq_name.NEXTVAL FROM DUAL.
Aurora PostgreSQL: same concept, use SELECT nextval('seq_name').
Key change: no more 'FROM DUAL'. Use SERIAL or GENERATED AS IDENTITY for auto-increment columns.
Migration tip: replace DUAL-based sequence calls with nextval() function calls.""",

    """Oracle DATE vs Aurora PostgreSQL TIMESTAMP
Oracle DATE stores date AND time (to the second) — a common gotcha.
Aurora PostgreSQL DATE stores date only. Use TIMESTAMP or TIMESTAMPTZ for date+time.
Migration tip: convert all Oracle DATE columns to TIMESTAMP in Aurora to avoid silent data truncation.""",

    """Oracle NVL vs Aurora PostgreSQL COALESCE
Oracle: NVL(expr, fallback) returns fallback when expr is NULL.
Aurora PostgreSQL: COALESCE(expr, fallback) — ANSI SQL standard, identical behaviour for two args.
COALESCE accepts multiple fallbacks; NVL accepts only two.
Migration tip: NVL(x, y) becomes COALESCE(x, y). No logic change needed.""",

    """Oracle ROWNUM vs Aurora PostgreSQL LIMIT / OFFSET
Oracle: WHERE ROWNUM <= 10 for top-N queries.
Aurora PostgreSQL: LIMIT 10. For pagination: LIMIT 10 OFFSET 20.
Oracle pagination requires nested subqueries with ROWNUM; Aurora uses clean LIMIT/OFFSET.
Window function ROW_NUMBER() OVER (ORDER BY col) works in both if you need ranked results.""",

    """Aurora PostgreSQL Connection Best Practices for Oracle DBAs
Use the cluster writer endpoint for writes; reader endpoint for read-only queries.
Aurora auto-scales storage and supports up to 15 read replicas.
Connection pooling: use RDS Proxy (managed) or PgBouncer (self-managed) to reduce connection overhead.
Max connections depend on instance class — r6g.large supports ~680, r6g.2xlarge supports ~1700.
Unlike Oracle dedicated server processes, Aurora uses a shared process model.""",
]


# ── Main ────────────────────────────────────────────────────────────────────────
if __name__ == "__main__":
    print("=" * 60)
    print("  RAG App: Amazon Bedrock + Aurora PostgreSQL (pgvector)")
    print("=" * 60)

    print("\n[1/3] Setting up Aurora pgvector...")
    setup_db()

    print("\n[2/3] Ingesting Oracle migration docs...")
    ingest(ORACLE_MIGRATION_DOCS)

    question = "What is the difference between Oracle sequences and Aurora PostgreSQL sequences?"
    print(f"\n[3/3] Asking: {question}\n")

    answer = ask(question)

    print("\n" + "─" * 60)
    print("Answer:")
    print("─" * 60)
    print(answer)
    print("─" * 60)