import sys
import requests
import logging
from langchain_community.document_loaders import PDFPlumberLoader
from langchain_experimental.text_splitter import SemanticChunker
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_ollama import OllamaLLM
from langchain_core.prompts import PromptTemplate
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain.chains.retrieval import create_retrieval_chain
from langchain_chroma import Chroma
from langchain_core.documents import Document
from bs4 import BeautifulSoup
import gradio as gr
from urllib.parse import urljoin

# Function to scrape content from a URL
def scrape_url(url):
    try:
        response = requests.get(url, timeout=3)
        
        # Check if the content is HTML
        content_type = response.headers.get('Content-Type', '').lower()
        if 'text/html' not in content_type:
            print(f"Skipping non-HTML content at {url}")
            return ""

        soup = BeautifulSoup(response.text, 'html.parser')
        return soup.get_text(strip=True)
    except requests.RequestException as e:
        print(f"Error scraping {url}: {e}")
        return ""

# Main website URL
main_url = "https://rocm.docs.amd.com/en/latest/"

# Scrape the main page
try:
    main_page = requests.get(main_url)
    main_soup = BeautifulSoup(main_page.text, 'html.parser')
except requests.RequestException as e:
    print(f"Error accessing main page: {e}")
    exit(1)

# Find all links on the main page
links = main_soup.find_all('a', href=True)

# Scrape content from the main page and linked pages
all_content = []
all_content.append(Document(page_content=scrape_url(main_url), metadata={"source": main_url}))
for link in links:
    full_url = urljoin(main_url, link['href'])
    if (full_url == "https://www.amd.com/"):
        print("detected main amd site-ignoring")
    elif (full_url == "https://www.amd.com/en/developer/resources/infinity-hub.html"):
        print("detected Infinity Hub-ignoring")
    elif full_url.startswith('http') and not full_url.lower().endswith('.pdf'):
        content = scrape_url(full_url)
        print(f"Scraping site: ", full_url)
        if content:
            all_content.append(Document(page_content=content, metadata={"source": full_url}))

llm = OllamaLLM(model="llama3.3:70b")

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
texts = text_splitter.split_documents(all_content)

embeddings = HuggingFaceEmbeddings()

vectorstore = Chroma.from_documents(texts, embeddings)

retriever = vectorstore.as_retriever()

prompt = PromptTemplate.from_template("""
1. Use the following pieces of context to answer the question related to AMD's ROCm product.
2. If you don't know the answer, just say that "I don't know" but don't make up an answer on your own.
3. Keep the answer crisp and limited to 3,4 sentences, and try to only use the context provided when answering.

Context: {context}

Question: {input}

Helpful Answer:""")

combine_docs_chain = create_stuff_documents_chain(llm, prompt)
qa = create_retrieval_chain(retriever, combine_docs_chain)

def respond(question, history=None):
    return qa.invoke({"input": question})["answer"]

def run_cli():
    print("\nAMD AI Assistant Ready! Type your questions. Type 'exit', 'quit' or 'bye' to stop.\n")
    while True:
        user_input = input("Prompt: ")
        if user_input.lower() in ["exit", "quit", "bye"]:
            print("Goodbye!")
            break
        answer = respond(user_input)
        print(f"AMD AI Assistant: {answer}\n")

def run_webui():
    gr.ChatInterface(
        respond,
        chatbot=gr.Chatbot(height=500),
        textbox=gr.Textbox(placeholder="Ask me question related to the awesomeness of ROCm and how it can revolutionize your AI workflows", container=False, scale=7),
        title="Rocm-bot",
        examples=["How can I install ROCm", "What installation methods exist for ROCm"],
        cache_examples=False,
        theme="Glass",
    ).launch(share=True, server_name="0.0.0.0")

if __name__ == "__main__":
    if "--webui" in sys.argv:
        run_webui()
    else:
        run_cli()

