import os
import streamlit as st
import requests
import json
from bs4 import BeautifulSoup
import re
from difflib import SequenceMatcher

# Fix Streamlit permission issue
os.environ['STREAMLIT_HOME'] = '/tmp/.streamlit'

API_URL = "https://jishnusetia-ollama-server.hf.space/api/generate"
MODEL = "llama3.2"

st.title("🌐 Chat with a Website (Ollama)")

# Store website content in memory
if "chunks" not in st.session_state:
    st.session_state.chunks = []
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# Helper: clean and chunk text
def clean_and_chunk(text, chunk_size=800):
    text = re.sub(r'\s+', ' ', text).strip()
    return [text[i:i+chunk_size] for i in range(0, len(text), chunk_size)]

# Helper: get top N relevant chunks
def get_relevant_chunks(query, chunks, top_n=3):
    scored = [(SequenceMatcher(None, query, chunk).ratio(), chunk) for chunk in chunks]
    scored.sort(key=lambda x: x[0], reverse=True)
    return [chunk for _, chunk in scored[:top_n]]

# Step 1: User inputs website
url = st.text_input("Enter website URL:")

if st.button("Fetch Website") and url:
    try:
        resp = requests.get(url, timeout=10)
        soup = BeautifulSoup(resp.text, 'html.parser')
        for script in soup(["script", "style"]):
            script.extract()
        text = soup.get_text()
        st.session_state.chunks = clean_and_chunk(text)
        st.success(f"Website fetched! {len(st.session_state.chunks)} chunks ready.")
    except Exception as e:
        st.error(f"Error fetching website: {e}")

# Step 2: Function to send query to Ollama
def send_message(message, context):
    prompt = f"Answer the following question based on this context:\n{context}\n\nQuestion: {message}"
    payload = {"model": MODEL, "prompt": prompt}
    response = requests.post(API_URL, json=payload, stream=True)

    if response.ok:
        reply = ""
        for line in response.iter_lines():
            if line:
                data = json.loads(line.decode('utf-8'))
                reply += data.get("response", "")
        return reply.strip()
    else:
        return "⚠️ Error: Failed to get response from Ollama API."

# Step 3: Chat input
user_input = st.text_input("Ask a question about the website:")

if st.button("Send Question") and user_input.strip():
    if not st.session_state.chunks:
        st.error("Please fetch a website first.")
    else:
        st.session_state.chat_history.append(("You", user_input))
        with st.spinner("Thinking..."):
            context = "\n".join(get_relevant_chunks(user_input, st.session_state.chunks))
            bot_reply = send_message(user_input, context)
        st.session_state.chat_history.append(("Ollama", bot_reply))

# Step 4: Display chat history
st.subheader("Chat History")
for speaker, message in st.session_state.chat_history:
    if speaker == "You":
        st.markdown(f"**You:** {message}")
    else:
        st.markdown(f"**🤖 Ollama:** {message}")