import os
import feedparser
from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.docstore.document import Document
import logging
from huggingface_hub import HfApi, login
import shutil
import rss_feeds
from datetime import datetime
import dateutil.parser
import hashlib
import re  # For cleaning HTML and whitespace

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Constants
MAX_ARTICLES_PER_FEED = 5
LOCAL_DB_DIR = "chroma_db"
RSS_FEEDS = rss_feeds.RSS_FEEDS
COLLECTION_NAME = "news_articles"

HF_API_TOKEN = os.getenv("DEMO_HF_API_TOKEN", "YOUR_HF_API_TOKEN")
REPO_ID = "broadfield-dev/news-rag-db"

# Initialize Hugging Face API
login(token=HF_API_TOKEN)
hf_api = HfApi()

# Initialize embedding model (global, reusable)
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

# Initialize vector DB with a specific collection name
vector_db = Chroma(
    persist_directory=LOCAL_DB_DIR,
    embedding_function=embedding_model,
    collection_name=COLLECTION_NAME
)

def clean_text(text):
    """Clean text by removing HTML tags and extra whitespace."""
    if not text or not isinstance(text, str):
        return ""
    # Remove HTML tags
    text = re.sub(r'<.*?>', '', text)
    # Normalize whitespace (remove extra spaces, newlines, tabs)
    text = ' '.join(text.split())
    return text.strip().lower()

def fetch_rss_feeds():
    articles = []
    seen_keys = set()
    for feed_url in RSS_FEEDS:
        try:
            logger.info(f"Fetching {feed_url}")
            feed = feedparser.parse(feed_url)
            if feed.bozo:
                logger.warning(f"Parse error for {feed_url}: {feed.bozo_exception}")
                continue
            article_count = 0
            for entry in feed.entries:
                if article_count >= MAX_ARTICLES_PER_FEED:
                    break
                title = entry.get("title", "No Title")
                link = entry.get("link", "")
                description = entry.get("summary", entry.get("description", ""))
                
                # Clean and normalize all text fields
                title = clean_text(title)
                link = clean_text(link)
                description = clean_text(description)

                # Try multiple date fields and parse flexibly
                published = "Unknown Date"
                for date_field in ["published", "updated", "created", "pubDate"]:  # Added "pubDate" for broader compatibility
                    if date_field in entry:
                        try:
                            parsed_date = dateutil.parser.parse(entry[date_field])
                            published = parsed_date.strftime("%Y-%m-%d %H:%M:%S")
                            break
                        except (ValueError, TypeError) as e:
                            logger.debug(f"Failed to parse {date_field} '{entry[date_field]}': {e}")
                            continue

                # Use a robust key for deduplication, including cleaned fields
                description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()  # Switched to SHA256 for better uniqueness
                key = f"{title}|{link}|{published}|{description_hash}"
                if key not in seen_keys:
                    seen_keys.add(key)
                    # Try multiple image sources
                    image = "svg"  # Default fallback
                    for img_source in [
                        lambda e: clean_text(e.get("media_content", [{}])[0].get("url")) if e.get("media_content") else "",
                        lambda e: clean_text(e.get("media_thumbnail", [{}])[0].get("url")) if e.get("media_thumbnail") else "",
                        lambda e: clean_text(e.get("enclosure", {}).get("url")) if e.get("enclosure") else "",
                        lambda e: clean_text(next((lnk.get("href") for lnk in e.get("links", []) if lnk.get("type", "").startswith("image")), "")),
                    ]:
                        try:
                            img = img_source(entry)
                            if img and img.strip():
                                image = img
                                break
                        except (IndexError, AttributeError, TypeError):
                            continue

                    articles.append({
                        "title": title,
                        "link": link,
                        "description": description,
                        "published": published,
                        "category": categorize_feed(feed_url),
                        "image": image,
                    })
                    article_count += 1
                else:
                    logger.debug(f"Duplicate article skipped in feed {feed_url}: {key}")
        except Exception as e:
            logger.error(f"Error fetching {feed_url}: {e}")
    logger.info(f"Total articles fetched: {len(articles)}")
    return articles

def categorize_feed(url):
    if "nature" in url.lower() or "science.org" in url.lower() or "arxiv.org" in url.lower() or "plos.org" in url.lower() or "annualreviews.org" in url.lower() or "journals.uchicago.edu" in url.lower() or "jneurosci.org" in url.lower() or "cell.com" in url.lower() or "nejm.org" in url.lower() or "lancet.com" in url.lower():
        return "Academic Papers"
    elif "reuters.com/business" in url.lower() or "bloomberg.com" in url.lower() or "ft.com" in url.lower() or "marketwatch.com" in url.lower() or "cnbc.com" in url.lower() or "foxbusiness.com" in url.lower() or "wsj.com" in url.lower() or "bworldonline.com" in url.lower() or "economist.com" in url.lower() or "forbes.com" in url.lower():
        return "Business"
    elif "investing.com" in url.lower() or "cnbc.com/market" in url.lower() or "marketwatch.com/market" in url.lower() or "fool.co.uk" in url.lower() or "zacks.com" in url.lower() or "seekingalpha.com" in url.lower() or "barrons.com" in url.lower() or "yahoofinance.com" in url.lower():
        return "Stocks & Markets"
    elif "whitehouse.gov" in url.lower() or "state.gov" in url.lower() or "commerce.gov" in url.lower() or "transportation.gov" in url.lower() or "ed.gov" in url.lower() or "dol.gov" in url.lower() or "justice.gov" in url.lower() or "federalreserve.gov" in url.lower() or "occ.gov" in url.lower() or "sec.gov" in url.lower() or "bls.gov" in url.lower() or "usda.gov" in url.lower() or "gao.gov" in url.lower() or "cbo.gov" in url.lower() or "fema.gov" in url.lower() or "defense.gov" in url.lower() or "hhs.gov" in url.lower() or "energy.gov" in url.lower() or "interior.gov" in url.lower():
        return "Federal Government"
    elif "weather.gov" in url.lower() or "metoffice.gov.uk" in url.lower() or "accuweather.com" in url.lower() or "weatherunderground.com" in url.lower() or "noaa.gov" in url.lower() or "wunderground.com" in url.lower() or "climate.gov" in url.lower() or "ecmwf.int" in url.lower() or "bom.gov.au" in url.lower():
        return "Weather"
    elif "data.worldbank.org" in url.lower() or "imf.org" in url.lower() or "un.org" in url.lower() or "oecd.org" in url.lower() or "statista.com" in url.lower() or "kff.org" in url.lower() or "who.int" in url.lower() or "cdc.gov" in url.lower() or "bea.gov" in url.lower() or "census.gov" in url.lower() or "fdic.gov" in url.lower():
        return "Data & Statistics"
    elif "nasa" in url.lower() or "spaceweatherlive" in url.lower() or "space" in url.lower() or "universetoday" in url.lower() or "skyandtelescope" in url.lower() or "esa" in url.lower():
        return "Space"
    elif "sciencedaily" in url.lower() or "quantamagazine" in url.lower() or "smithsonianmag" in url.lower() or "popsci" in url.lower() or "discovermagazine" in url.lower() or "scientificamerican" in url.lower() or "newscientist" in url.lower() or "livescience" in url.lower() or "atlasobscura" in url.lower():
        return "Science"
    elif "wired" in url.lower() or "techcrunch" in url.lower() or "arstechnica" in url.lower() or "gizmodo" in url.lower() or "theverge" in url.lower():
        return "Tech"
    elif "horoscope" in url.lower() or "astrostyle" in url.lower():
        return "Astrology"
    elif "cnn_allpolitics" in url.lower() or "bbci.co.uk/news/politics" in url.lower() or "reuters.com/arc/outboundfeeds/newsletter-politics" in url.lower() or "politico.com/rss/politics" in url.lower() or "thehill" in url.lower():
        return "Politics"
    elif "weather" in url.lower() or "swpc.noaa.gov" in url.lower() or "foxweather" in url.lower():
        return "Earth Weather"
    elif "vogue" in url.lower():
        return "Lifestyle"
    elif "phys.org" in url.lower() or "aps.org" in url.lower() or "physicsworld" in url.lower():
        return "Physics"
    return "Uncategorized"

def process_and_store_articles(articles):
    documents = []
    existing_ids = set(vector_db.get()["ids"])  # Get existing document IDs to avoid duplicates
    for article in articles:
        try:
            # Clean and normalize all fields
            title = clean_text(article["title"])
            link = clean_text(article["link"])
            description = clean_text(article["description"])
            published = article["published"]
            description_hash = hashlib.sha256(description.encode('utf-8')).hexdigest()
            doc_id = f"{title}|{link}|{published}|{description_hash}"
            if doc_id in existing_ids:
                logger.debug(f"Skipping duplicate in DB: {doc_id}")
                continue
            metadata = {
                "title": article["title"],
                "link": article["link"],
                "original_description": article["description"],
                "published": article["published"],
                "category": article["category"],
                "image": article["image"],
            }
            doc = Document(page_content=description, metadata=metadata, id=doc_id)
            documents.append(doc)
        except Exception as e:
            logger.error(f"Error processing article {article['title']}: {e}")
    
    if documents:
        try:
            vector_db.add_documents(documents)
            vector_db.persist()  # Explicitly persist changes
            logger.info(f"Added {len(documents)} new articles to DB")
        except Exception as e:
            logger.error(f"Error storing articles: {e}")

def download_from_hf_hub():
    # Only download if the local DB doesn’t exist (initial setup)
    if not os.path.exists(LOCAL_DB_DIR):
        try:
            hf_api.create_repo(repo_id=REPO_ID, repo_type="dataset", exist_ok=True, token=HF_API_TOKEN)
            logger.info(f"Downloading Chroma DB from {REPO_ID}...")
            hf_api.download_repo(repo_id=REPO_ID, repo_type="dataset", local_dir=LOCAL_DB_DIR, token=HF_API_TOKEN)
        except Exception as e:
            logger.error(f"Error downloading from Hugging Face Hub: {e}")
            raise
    else:
        logger.info("Local Chroma DB already exists, skipping download.")

def upload_to_hf_hub():
    if os.path.exists(LOCAL_DB_DIR):
        try:
            logger.info(f"Uploading updated Chroma DB to {REPO_ID}...")
            for root, _, files in os.walk(LOCAL_DB_DIR):
                for file in files:
                    local_path = os.path.join(root, file)
                    remote_path = os.path.relpath(local_path, LOCAL_DB_DIR)
                    hf_api.upload_file(
                        path_or_fileobj=local_path,
                        path_in_repo=remote_path,
                        repo_id=REPO_ID,
                        repo_type="dataset",
                        token=HF_API_TOKEN
                    )
            logger.info(f"Database uploaded to: {REPO_ID}")
        except Exception as e:
            logger.error(f"Error uploading to Hugging Face Hub: {e}")
            raise

if __name__ == "__main__":
    articles = fetch_rss_feeds()
    process_and_store_articles(articles)
    upload_to_hf_hub()