import os
import requests
import gradio as gr

# ---------------- RAG DOCUMENT SETUP ---------------- #

K8S_DOC_URLS = {
    "pods": "https://kubernetes.io/docs/concepts/workloads/pods/",
    "deployments": "https://kubernetes.io/docs/concepts/workloads/controllers/deployment/",
    "services": "https://kubernetes.io/docs/concepts/services-networking/service/",
    "namespaces": "https://kubernetes.io/docs/concepts/overview/working-with-objects/namespaces/",
    "nodes": "https://kubernetes.io/docs/concepts/architecture/nodes/",
    "statefulsets": "https://kubernetes.io/docs/concepts/workloads/controllers/statefulset/",
    "rbac": "https://kubernetes.io/docs/reference/access-authn-authz/rbac/",
    "persistent-volumes": "https://kubernetes.io/docs/concepts/storage/persistent-volumes/",
    "ingress": "https://kubernetes.io/docs/concepts/services-networking/ingress/",
    "autoscaling": "https://kubernetes.io/docs/tasks/run-application/horizontal-pod-autoscale/"
}

def fetch_doc(url):
    try:
        response = requests.get(url, timeout=10)
        if response.status_code == 200:
            return response.text
    except:
        return ""
    return ""

DOCUMENTS = [
    {"doc": name, "url": url, "text": fetch_doc(url)}
    for name, url in K8S_DOC_URLS.items()
]

def search_docs(query, top_k=3):
    query = query.lower()
    matches = []
    for doc in DOCUMENTS:
        text = doc["text"].lower()
        if query in text:
            snippet_start = text.index(query)
            snippet_end = snippet_start + 350
            snippet = doc["text"][snippet_start:snippet_end].replace("\n", " ")
            matches.append((snippet, doc["url"], doc["doc"]))
    return matches[:top_k]


# --------------- LLM CALL (OpenRouter) ---------------- #

def call_llm(prompt):
    url = "https://openrouter.ai/api/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {os.getenv('OPENROUTER_API_KEY')}",
        "HTTP-Referer": "https://huggingface.co/",
        "X-Title": "Kubernetes RAG Assistant"
    }

    data = {
        "model": "meta-llama/llama-3.1-8b-instruct",
        "messages": [{"role": "user", "content": prompt}],
        "max_tokens": 350
    }

    res = requests.post(url, json=data, headers=headers)
    out = res.json()

    if "choices" in out:
        return out["choices"][0]["message"]["content"]
    print("DEBUG LLM Error:", out)
    return "⚠ Model error. Try again."


# ----------- RAG + Prompt Construction ---------------- #

def build_answer(query):
    results = search_docs(query)
    context = ""
    citations = []

    for i, (snippet, url, doc) in enumerate(results, start=1):
        label = f"[{i}]"
        context += f"{label}: {snippet}\n\n"
        citations.append(f"{label} → {url}")

    prompt = f"""
    Use the context below to answer the question clearly.
    Add citations like [1], [2] at the end of sentences.

    Context:
    {context}

    Question: {query}
    """

    answer = call_llm(prompt)
    citations_text = "\n".join(citations) or "No sources found."

    return answer, citations_text


# ---------------------- UI --------------------------- #

custom_css = """
.source-box {
    font-size: 14px;
    background: #1b2733;
    padding: 10px;
    border-radius: 8px;
    color: #c9e2ff;
    border: 1px solid #4a90e2;
}
"""

with gr.Blocks(css=custom_css, theme="soft") as app:

    gr.HTML("""
    <h1 style='color:#326ce5; text-align:center;'>☸️ Kubernetes RAG Assistant</h1>
    <p style='text-align:center; font-size:17px; color:#ddd;'>Ask any Kubernetes question and get answers with docs citations 📌</p>
    """)

    question = gr.Textbox(label="Ask a Kubernetes Question:", placeholder="e.g., What is RBAC in Kubernetes?")

    answer = gr.Markdown(label="Answer")
    sources = gr.Markdown(label="Sources", elem_classes=["source-box"])

    submit = gr.Button("Ask ☸️")

    submit.click(build_answer, inputs=question, outputs=[answer, sources])

app.launch()