import os
import psycopg2
from openai import OpenAI
from dotenv import load_dotenv

# Load OpenAI API key from .env file
load_dotenv()
client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# PostgreSQL connection config
PG_CONN_INFO = {
    "host": "localhost",
    "port": 5432,
    "dbname": "pgdevai",
    "user": "garyevans",
    "password": "*********"
}

# Embed user query using OpenAI
def embed_query(text):
    response = client.embeddings.create(
        input=text,
        model="text-embedding-3-small"
    )
    return response.data[0].embedding

# Retrieve top-k semantically similar chunks from manual_chunks
def retrieve_chunks(query_vector, top_nbr=15):
    with psycopg2.connect(**PG_CONN_INFO) as conn:
        with conn.cursor() as cur:
            cur.execute("""
                SELECT section_title, chunk, method, section_anchor
                FROM production.manual_chunks
                WHERE embedding IS NOT NULL
                ORDER BY embedding <-> %s::vector
                LIMIT %s;
            """, (query_vector, top_nbr))
            return cur.fetchall()

# Format retrieved chunks and generate a grounded response
def generate_response(user_query, chunks):
    context = "\n\n".join([
        f"[{i+1}] Section: {title}\nMethod: {method}\n{chunk.strip()}"
        for i, (title, chunk, method, anchor) in enumerate(chunks)
    ])

    prompt = f"""
You are a technical assistant answering a user question using internal technical manuals.

User's Question:
"{user_query}"

Below are retrieved documentation chunks from a semantic search:

{context}

INSTRUCTIONS:
- Read the retrieved content.
- Determine which sections answer the question directly.
- Extract only the information relevant to the question.
- Synthesize a grounded and helpful answer.
- Do NOT repeat full paragraphs.
- Use bullet points or summarise as needed.
- If the content is incomplete or unclear, say so.

Answer:
"""

    response = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system", "content": "You are a precise and helpful technical assistant who explains only what is supported by the provided documentation."},
            {"role": "user", "content": prompt}
        ],
        temperature=0.3
    )
    return response.choices[0].message.content

# ---- MAIN ----
if __name__ == "__main__":
    query = input("Enter your question: ").strip()
    embedded = embed_query(query)
    results = retrieve_chunks(embedded)
    answer = generate_response(query, results)
    print(answer)