import os
import requests
import json
from bs4 import BeautifulSoup
import streamlit as st
from huggingface_hub import InferenceClient
from datasets import load_dataset
from duckduckgo_search import DDGS
import yaml  # Importing YAML module

# Load the dataset
ds = load_dataset("nvidia/HelpSteer2")

# Initialize Hugging Face client
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

# Load premade prompts from YAML with error handling
def load_premade_prompts(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return yaml.safe_load(file)  # Load YAML instead of JSON
    except yaml.YAMLError as e:
        st.error(f"Error decoding YAML: {e}")
        return {}
    except Exception as e:
        st.error(f"An error occurred: {e}")
        return {}

# Load prompts
prompts = load_premade_prompts("J:/project_relapse/agent_jobs/templates/premade_prompts.yaml").get("categories", {})

# Specify the output file for scraped content
scraped_content_file = "templates/scraped_content.txt"

# Function to search DuckDuckGo
def search_duckduckgo(query):
    results = DDGS().text(query, max_results=5)
    return results

# Function to scrape text and URLs from a page
def scrape_all_text_and_urls(url, is_email_content=False):
    try:
        response = requests.get(url)
        response.raise_for_status()  # Raise an HTTPError for 4xx and 5xx status codes
        soup = BeautifulSoup(response.text, 'html.parser')

        if is_email_content:
            email_content = soup.find_all(text=True)
            clean_text = ' '.join(email_content).replace('\n', ' ').strip()
            return clean_text

        all_text = soup.get_text(separator='\n', strip=True)
        return all_text
    except Exception as e:
        st.error(f"Error scraping the website '{url}': {str(e)}")
        return None

# Function to chunk content
def chunk_content(content, max_chunk_size):
    chunks = []
    words = content.split()
    for i in range(0, len(words), max_chunk_size):
        chunks.append(' '.join(words[i:i + max_chunk_size]))
    return chunks

# Function to summarize content with Hugging Face model
def summarize_content_with_huggingface(content, prompt, max_tokens=4096):
    chunks = chunk_content(content, 1500)  # Adjust chunk size based on model limits
    summary = ""
    summary_output = st.empty()  # Create a placeholder for the summary output

    for chunk in chunks:
        messages = [{"role": "system", "content": prompt}]
        messages.append({"role": "user", "content": chunk})

        for response in client.chat_completion(messages, stream=True, max_tokens=max_tokens):
            delta_content = response.choices[0].delta.content
            summary += delta_content
            
            # Update the placeholder with the current summary
            summary_output.markdown(f"### Summarized Content:\n{summary}")

    return summary

# Streamlit app layout
st.title("DuckDuckGo Search and Content Summarizer")
st.write("Enter a search term to find relevant links and summarize their content.")

# Dropdown for categories
selected_category = st.selectbox("Select a category:", list(prompts.keys()))
selected_prompt = prompts[selected_category]["prompt"] if selected_category in prompts else ""

# Dropdown for premade searches within the selected category
selected_search = st.selectbox("Select a search type:", prompts[selected_category]["searches"] if selected_category in prompts else [])

# User input for search term, prefilled with the selected prompt if needed
query = st.text_input("Search Term (e.g., 'Italian restaurants')", value=selected_search, key="search_input")

# Input for direct URL scraping
url_input = st.text_input("Or provide a URL to scrape directly:", placeholder="http://example.com")

# Search and process button
if st.button("Search, Scrape & Summarize"):
    scraped_content = ""

    if url_input:
        is_email_content = selected_category == "Emails"
        scraped_content = scrape_all_text_and_urls(url_input, is_email_content)
    elif query:
        results = search_duckduckgo(query)
        if results:
            st.write("### Search Results:")
            urls = []
            for idx, result in enumerate(results, start=1):
                st.write(f"{idx}. **{result['title']}**")
                st.write(f"[Link]({result['href']})")
                st.write(result['body'])
                urls.append(result['href'])

            for url in urls:
                is_email_content = selected_category == "Emails"
                scraped_content = scrape_all_text_and_urls(url, is_email_content)
                if scraped_content:
                    break

            if not scraped_content:
                st.warning("Failed to scrape all provided links.")
        else:
            st.write("No results found.")
    else:
        st.write("Please enter a search query or a URL.")

    if scraped_content:
        with open(scraped_content_file, 'w', encoding='utf-8') as file:
            file.write(scraped_content)
        st.success("Content scraped successfully and saved to output file.")

        st.write("### Scraped Content:")
        st.text_area("Scraped Content:", value=scraped_content, height=300, disabled=True)

        summary = summarize_content_with_huggingface(scraped_content, selected_prompt)
        st.write("### Summarized Content:")
        st.text_area("Summarized Content:", value=summary, height=300, disabled=True)

# Display filtered URLs (if any, modify as needed)
st.write("### Filtered URLs:")
if 'filtered_urls' in st.session_state:
    for url in st.session_state['filtered_urls']:
        st.write(f"- {url}")
else:
    st.write("No filtered URLs.")


can you tell me if this saves the generated response to a file?