"""
Module for sentiment analysis using FinBERT Indian news models.

First run downloads the model (~400MB). Subsequent runs use cached version.
Run on CPU — no GPU required for this model size.
"""

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Using the industry standard ProsusAI/finbert because kdave/FineTuned_Finbert 
# is missing model weights on the Hugging Face hub.
MODEL_NAME = "ProsusAI/finbert"

try:
    import streamlit as st
    cache_decorator = st.cache_resource(show_spinner="Loading FinBERT Model...")
except ImportError:
    from functools import lru_cache
    cache_decorator = lru_cache(maxsize=1)

@cache_decorator
def load_sentiment_model():
    """
    Loads the tokenizer and model from Hugging Face.
    Caches the model after the first load to avoid reloading.
    """
    print(f"Loading sentiment model '{MODEL_NAME}'... (This may take a moment on first run)")
    _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    _model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
    _model.eval()  # Set to evaluation mode
    _model.to("cpu")
    print("Sentiment model loaded successfully!")
        
    return _tokenizer, _model

def score_headline(headline_text):
    """
    Scores a single news headline string.
    Returns dict: {"label": "positive/negative/neutral", "score": confidence_score}
    Handles empty or None input gracefully.
    """
    if not headline_text or not str(headline_text).strip():
        return {"label": "neutral", "score": 0.0}
        
    tokenizer, model = load_sentiment_model()
    
    inputs = tokenizer(headline_text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    
    with torch.no_grad():
        outputs = model(**inputs)
        
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    top_prob, top_idx = torch.max(probs, dim=-1)
    
    top_prob = top_prob.item()
    top_idx = top_idx.item()
    
    raw_label = model.config.id2label.get(top_idx, "neutral").lower()
    
    if "pos" in raw_label:
        label = "positive"
    elif "neg" in raw_label:
        label = "negative"
    else:
        label = "neutral"
        
    return {"label": label, "score": top_prob}

def score_portfolio_news(news_list):
    """
    Scores a list of news dictionaries.
    Input format: [{"symbol": "TCS", "headline": "TCS beats Q4 estimates"}]
    Returns enriched list with sentiment and score added.
    """
    if not news_list:
        return []
        
    enriched_news = []
    for item in news_list:
        headline = item.get("headline", "")
        sentiment_data = score_headline(headline)
        
        enriched_item = item.copy()
        enriched_item["sentiment"] = sentiment_data["label"]
        enriched_item["score"] = round(sentiment_data["score"], 4)
        enriched_news.append(enriched_item)
        
    return enriched_news

def get_sentiment_emoji(label):
    """
    Returns emoji for UI based on sentiment label.
    """
    lbl = str(label).lower()
    if "pos" in lbl:
        return "🟢"
    elif "neg" in lbl:
        return "🔴"
    return "🟡"

if __name__ == "__main__":
    import sys
    if sys.platform == 'win32':
        sys.stdout.reconfigure(encoding='utf-8')
        
    sample_news = [
        {"symbol": "RELIANCE", "headline": "Reliance Industries announces massive ₹75,000 crore investment in green energy, stock surges."},
        {"symbol": "TCS", "headline": "TCS Q4 margins contract due to wage hikes and macroeconomic headwinds in the US."},
        {"symbol": "HDFCBANK", "headline": "HDFC Bank holds interest rates steady, analysts predict stable growth for the quarter."}
    ]
    
    print("Testing FinBERT Sentiment Module...\n")
    results = score_portfolio_news(sample_news)
    
    for res in results:
        symbol = res["symbol"]
        headline = res["headline"]
        sentiment = res["sentiment"]
        score = res["score"]
        emoji = get_sentiment_emoji(sentiment)
        
        print(f"[{symbol}] {headline}")
        print(f"--> Sentiment: {emoji} {sentiment.upper()} (Confidence: {score:.2f})\n")