import streamlit as st
import feedparser
import requests
import json
import os
from bs4 import BeautifulSoup
from transformers import MarianMTModel, MarianTokenizer
import whisper
from streamlit_audio_recorder import st_audiorec

st.set_page_config(page_title="📍 BharatPulse - Local News & Sentiment", layout="wide")

# ------------------------------
# User Auth Functions
# ------------------------------
USERS_FILE = "users.json"

def load_users():
    if not os.path.exists(USERS_FILE):
        return {}
    with open(USERS_FILE, "r") as f:
        return json.load(f)

def save_users(users):
    with open(USERS_FILE, "w") as f:
        json.dump(users, f, indent=2)

def register_user(username, email, password, location):
    users = load_users()
    if email in users:
        return False, "User already exists."
    users[email] = {
        "username": username,
        "email": email,
        "password": password,
        "location": location
    }
    save_users(users)
    return True, "Registration successful."

def login_user(email, password):
    users = load_users()
    user = users.get(email)
    if user and user["password"] == password:
        return True, user
    return False, "Invalid credentials."

# ------------------------------
# Whisper Model
# ------------------------------
@st.cache_resource
def load_whisper_model():
    return whisper.load_model("base")

def transcribe_audio(audio_bytes):
    path = "/tmp/voice.wav"
    with open(path, "wb") as f:
        f.write(audio_bytes)
    model = load_whisper_model()
    result = model.transcribe(path, language="te")
    return result["text"]

# ------------------------------
# Translation Model
# ------------------------------
@st.cache_resource
def load_translation_model():
    model_name = "Helsinki-NLP/opus-mt-en-te"
    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model

def translate_to_telugu(text):
    tokenizer, model = load_translation_model()
    tokens = tokenizer.prepare_seq2seq_batch([text], return_tensors="pt", truncation=True, max_length=512)
    translated = model.generate(**tokens)
    return tokenizer.decode(translated[0], skip_special_tokens=True)

# ------------------------------
# Extract Article Summary
# ------------------------------
def extract_article_content(url):
    try:
        response = requests.get(url, timeout=5)
        soup = BeautifulSoup(response.text, "html.parser")
        paragraphs = soup.find_all('p')
        content = ' '.join([p.get_text() for p in paragraphs[:5]])
        return content
    except Exception as e:
        return str(e)

# ------------------------------
# Auth UI
# ------------------------------
if "logged_in" not in st.session_state:
    st.session_state.logged_in = False

if not st.session_state.logged_in:
    st.title("🔐 Login to BharatPulse")

    tab1, tab2 = st.tabs(["🔑 Login", "📝 Register"])

    with tab1:
        email = st.text_input("Email")
        password = st.text_input("Password", type="password")
        if st.button("Login"):
            success, user_or_msg = login_user(email, password)
            if success:
                st.session_state.logged_in = True
                st.session_state.user = user_or_msg
                st.rerun()
            else:
                st.error(user_or_msg)

    with tab2:
        new_username = st.text_input("Username")
        new_email = st.text_input("New Email")
        new_password = st.text_input("New Password", type="password")
        new_location = st.text_input("Location")
        if st.button("Register"):
            success, message = register_user(new_username, new_email, new_password, new_location)
            if success:
                st.success(message)
            else:
                st.error(message)
    st.stop()

# ------------------------------
# Main News UI
# ------------------------------
st.sidebar.success(f"Logged in as: {st.session_state.user['username']}")
if st.sidebar.button("🚪 Logout"):
    st.session_state.logged_in = False
    st.session_state.user = None
    st.rerun()

st.title("📍 BharatPulse")
tabs = st.tabs(["📰 Headlines", "🔍 Search by City", "🎤 Voice Search"])

feeds = {
    "Sakshi - Medak": "https://www.sakshi.com/rss/district-news/medak.xml",
    "Eenadu - AP": "https://www.eenadu.net/rss/andhra-pradesh.xml"
}

# ------------------------------
# 📰 Headlines Tab
# ------------------------------
with tabs[0]:
    st.header("📰 Local Telugu News")
    for name, url in feeds.items():
        st.subheader(name)
        feed = feedparser.parse(url)
        for entry in feed.entries[:3]:
            st.markdown(f"### {entry.title}")
            st.write(entry.published)
            st.markdown(f"[Read Full Article]({entry.link})")
            summary = extract_article_content(entry.link)
            st.markdown("**English Summary:** " + summary)
            telugu_summary = translate_to_telugu(summary)
            st.markdown("**తెలుగు అనువాదం:** " + telugu_summary)
            st.markdown("---")

# ------------------------------
# 🔍 Search by City
# ------------------------------
with tabs[1]:
    st.header("🔎 Search by City / District")
    city = st.text_input("Enter your city/district name in Telugu or English")

    if city:
        st.success(f"Showing news for: {city}")
        for name, url in feeds.items():
            if city.lower() in name.lower():
                st.subheader(name)
                feed = feedparser.parse(url)
                for entry in feed.entries[:2]:
                    st.markdown(f"**🗞️ {entry.title}**")
                    summary = extract_article_content(entry.link)
                    telugu_summary = translate_to_telugu(summary)
                    st.markdown("**తెలుగు అనువాదం:** " + telugu_summary)

# ------------------------------
# 🎤 Voice Search Tab
# ------------------------------
with tabs[2]:
    st.header("🎤 Search by Telugu Voice")
    audio = st_audiorec()
    if audio:
        transcribed = transcribe_audio(audio)
        st.success(f"Voice Input Recognized: {transcribed}")
        city = transcribed
        for name, url in feeds.items():
            if city.lower() in name.lower():
                st.subheader(name)
                feed = feedparser.parse(url)
                for entry in feed.entries[:2]:
                    st.markdown(f"**🗞️ {entry.title}**")
                    summary = extract_article_content(entry.link)
                    telugu_summary = translate_to_telugu(summary)
                    st.markdown("**తెలుగు అనువాదం:** " + telugu_summary)