streamlit_app.py

import streamlit as st
import pandas as pd
from main import parse_query
from agents import fetch_data
from state import TraderState
from graph import compiled_graph
from langchain_openai import ChatOpenAI
from langchain_groq import ChatGroq
import os
import json
from fpdf import FPDF
import tempfile
import base64
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import pytz
import re
from pathlib import Path

# LLM Configuration
llm_parser = ChatGroq(model="moonshotai/kimi-k2-instruct-0905", api_key=os.getenv("GROQ_API_KEY"), temperature=0.0)

# Page Configuration
st.set_page_config(page_title="AI Trader Analyzer", page_icon="📈", layout="wide")

def apply_custom_styles():
    """Apply custom CSS styles"""
    current_dir = Path(__file__).parent
    css_file = current_dir / "style.css"
    
    try:
        with open(css_file, "r") as f:
            css_content = f.read()
        st.markdown(f"<style>{css_content}</style>", unsafe_allow_html=True)
    except FileNotFoundError as e:
        st.error(f"File not found: {e}")
    except Exception as e:
        st.error(f"Error loading styles: {e}")

apply_custom_styles()

def sanitize_text(text: str) -> str:
    replacements = {"—": "-", """: '"', """: '"', "'": "'", "'": "'", "…": "..."}
    for old, new in replacements.items():
        text = text.replace(old, new)
    return ''.join(c for c in text if ord(c) < 128)

def generate_alert_message_static(state: TraderState, symbol: str, action: str, horizon: str) -> str:
    confidence = state['confidence']
    raw_data = state.get('raw_data', {}) or {}
    indicators = raw_data.get("indicators", {})
    prediction = raw_data.get("prediction", "No prediction available.")
    holding = raw_data.get("holding", {})
    claim = state.get('claim', 'No claim.')
    skepticism = state.get('skepticism', 'No skepticism.')
    features = raw_data.get("features", "")
    
    rsi = indicators.get('rsi', 50)
    macd = indicators.get('macd', 0)
    current_price = indicators.get('current_price', 0)
    currency = "Rs." if symbol.endswith(('.NS', '.BO')) else "$"
    
    volatility = 0
    if current_price > 0:
        bb_upper = indicators.get('bb_upper', current_price * 1.1)
        bb_lower = indicators.get('bb_lower', current_price * 0.9)
        volatility = (bb_upper - bb_lower) / current_price * 100
    
    risk_adjustment = min(20, volatility / 5) if volatility > 0 else 0
    
    if holding:
        pnl_adjustment = 10 if holding.get('holding_return', 0) > 5 else -10 if holding.get('holding_return', 0) < -5 else 0
        allocation = max(0, min(100, confidence - risk_adjustment + pnl_adjustment))
    else:
        allocation = max(0, min(100, confidence - risk_adjustment))
    
    if indicators.get("error") == "No data available":
        return f"No Data Available\nSorry, we couldn't fetch reliable data for {symbol}."
    
    horizon_desc = {"intraday": "same-day trading (hours)", "scalping": "ultra-short (minutes)", 
                    "swing": "1-4 weeks", "momentum": "2-4 weeks", "long_term": "months to years"}
    
    query_date = state.get('query_date')
    target_date = calculate_target_date(query_date, horizon)
    
    best_date_prompt = f"""Based on prediction: '{prediction}', RSI {rsi:.1f}, MACD {macd:.2f}.
    Suggest best date/time to {action} {symbol} for {horizon}. Be specific and include % gain/loss."""
    try:
        best_date = llm_parser.invoke(best_date_prompt).content.strip()
    except:
        best_date = f"Based on trends, {action} within {horizon_desc.get(horizon, horizon)}."
    
    why_prompt = f"""Based on RSI {rsi:.1f}, MACD {macd:.2f}, explain signal for {action} {symbol}. Use layman terms."""
    try:
        why = llm_parser.invoke(why_prompt).content.strip()
    except:
        why = f"Conditions are {'favorable' if confidence > 60 else 'mixed' if confidence > 40 else 'unfavorable'}."
    
    if confidence >= 70:
        color = " Green: 'Yes, Go Ahead!'"
        should_action = f"Yes, {action} now. Allocate {allocation:.0f}%."
    elif confidence >= 40:
        color = " Yellow: 'Wait and Watch'"
        should_action = f"Maybe, monitor closely. Allocate {allocation:.0f}% cautiously."
    else:
        color = " Red: 'No, Stop!'"
        should_action = f"No, avoid {action}. Wait for better conditions."
    
    holding_summary = ""
    if holding:
        h_return = holding['holding_return']
        holding_summary = f"\n**Holding**: {holding['days_held']} days, {currency}{holding['purchase_price']:.2f} → {currency}{holding.get('current_price', current_price):.2f}, {h_return:.1f}% ({holding['pnl']})."
        advice = "lock in gains" if h_return > 0 else "wait for recovery" if h_return < 0 else "exit at breakeven"
        should_action += f" {advice}."

    return f"""--- AI Trader Analyzer Report ---
Symbol: {symbol} | **Price**: {currency}{current_price:,.4f}
Action: {action.capitalize()} ({horizon_desc.get(horizon, horizon)}) | Target: {target_date}
Signal: {color} | Confidence: {confidence}%
Metrics: RSI {rsi:.1f}, MACD {macd:.2f}{holding_summary}
**Should I {action.capitalize()}?** {should_action}
**Best Time**: {best_date}
**Why?** {why}
**Prediction**: {prediction}
**Recommendation**: Allocate {allocation:.0f}%. Not financial advice."""

def classify_and_parse_query(user_input: str):
    prompt = f"""
    Analyze this user query: "{user_input}".
    
    TRADING HORIZONS (CRITICAL):
    - "intraday": same day buy/sell (hours)
    - "scalping": ultra-short (minutes)
    - "swing": 1-4 weeks (e.g., "next week", "middle of next month")
    - "momentum": 2-4 weeks trend-following
    - "long_term": months to years (e.g., "bought 1 month ago")
    
    Extract:
    - Is this a comparison query? (mentions 'vs', 'versus', 'compare', 'or', two stocks) Answer YES or NO
    - Stock symbols (e.g., "ICICIBANK.NS"). Infer with appropriate suffixes (.NS for Indian stocks)
    - action: "buy" or "sell" (default "buy")
    - horizon: match to above definitions
    - date: Extract timeframe (e.g., "middle of next month", "next week", "today", "tomorrow")
    - holding_period: If mentions "bought X ago", extract as "X unit ago", else null
    
    CRITICAL RULES:
    - "middle of next month" → horizon = "swing", date = "middle of next month"
    - "next week" → horizon = "swing", date = "next week"
    - "now" without future mention → horizon = "intraday", date = "today"
    - "bought X ago" → horizon = "long_term", holding_period = "X unit ago"
    
    Respond ONLY in JSON: {{"is_comparison": true/false, "symbols": ["ICICIBANK.NS", "HDFCBANK.NS"], "other_details": {{"action": "buy", "horizon": "swing", "date": "middle of next month", "holding_period": null}}}}
    """
    try:
        response = llm_parser.invoke(prompt).content.strip()
        parsed = json.loads(response)
        return parsed
    except:
        symbol, date, action, horizon, holding_period = parse_query(user_input)
        return {"is_comparison": False, "symbols": [symbol] if symbol else [], "other_details": {"action": action, "horizon": horizon, "date": date, "holding_period": holding_period}}

def calculate_target_date(date_str: str, horizon: str) -> str:
    ist = pytz.timezone('Asia/Kolkata')
    now_ist = datetime.now(ist)
    
    if horizon == 'scalping':
        return (now_ist + timedelta(minutes=10)).strftime('%d %b %Y, %I:%M %p IST')
    
    if not date_str or date_str in ['current', 'None']:
        days = {'intraday': 0, 'swing': 14, 'momentum': 21}.get(horizon, 0)
        return (now_ist + timedelta(days=days)).strftime('%d %b %Y')
    
    date_lower = str(date_str).lower()
    
    if 'middle of next month' in date_lower:
        next_month = now_ist.replace(year=now_ist.year + (1 if now_ist.month == 12 else 0),
                                       month=1 if now_ist.month == 12 else now_ist.month + 1, day=15)
        return next_month.strftime('%d %b %Y')
    elif 'next week' in date_lower:
        return (now_ist + timedelta(days=7)).strftime('%d %b %Y')
    elif 'tomorrow' in date_lower:
        return (now_ist + timedelta(days=1)).strftime('%d %b %Y')
    elif 'today' in date_lower:
        return now_ist.strftime('%d %b %Y')
    
    match = re.search(r'after\s+(\d+)\s*(week|month)', date_lower)
    if match:
        num, unit = int(match.group(1)), match.group(2)
        days = num * (7 if unit == "week" else 30)
        return (now_ist + timedelta(days=days)).strftime('%d %b %Y')
    
    try:
        return datetime.strptime(date_str, '%Y-%m-%d').strftime('%d %b %Y')
    except:
        return now_ist.strftime('%d %b %Y')

def get_horizon_description(horizon: str) -> str:
    return {"intraday": "Same-day (hours)", "scalping": "Ultra-short (minutes)", 
            "swing": "1-4 weeks", "momentum": "2-4 weeks", "long_term": "Months to years"}.get(horizon, horizon)

def process_normal_query(symbol, action, horizon, date, holding_period):
    hist = fetch_data(symbol, horizon)
    trader_state = TraderState(input_symbol=symbol, query_date=date, action=action, horizon=horizon, 
                                holding_period=holding_period, raw_data={'hist': hist}, claim=None, 
                                skepticism=None, confidence=50, iterations=0, stop=False, alert_message=None)
    final_state = compiled_graph.invoke(trader_state)
    alert_message = generate_alert_message_static(final_state, symbol, action, horizon)
    
    chart_data = []
    if not hist.empty:
        for idx, row in hist[['Close']].iterrows():
            chart_data.append({"date": idx.strftime('%Y-%m-%d'), "price": float(row['Close'])})
    
    return alert_message, chart_data, final_state

def process_comparison_query(symbol1, symbol2, action, horizon, date, holding_period):
    results = {}
    charts = {}
    states = {}
    for symbol in [symbol1, symbol2]:
        alert, chart_data, final_state = process_normal_query(symbol, action, horizon, date, holding_period)
        results[symbol] = alert
        charts[symbol] = chart_data
        states[symbol] = final_state
    
    pdf = FPDF()
    pdf.set_margins(left=6.35, top=12.7, right=6.35)
    pdf.set_auto_page_break(auto=True, margin=12.7)
    pdf.add_page()
    
    effective_width = pdf.w - 6.35 - 6.35
    
    pdf.set_font("Arial", 'B', size=16)
    pdf.cell(effective_width, 10, txt=sanitize_text(f"Stock Comparison Report: {symbol1} vs {symbol2}"), ln=True, align='C')
    pdf.ln(8)
    
    target_date = calculate_target_date(date, horizon)
    horizon_desc = get_horizon_description(horizon)
    
    pdf.set_font("Arial", 'B', size=13)
    pdf.cell(effective_width, 8, txt=sanitize_text("Overview"), ln=True)
    pdf.set_font("Arial", size=11)
    
    overview_text = f"Action: {action.capitalize()} | Horizon: {horizon.replace('_', ' ').title()} ({horizon_desc}) | Target Date: {target_date} | Holding Period: {holding_period or 'N/A'}"
    pdf.multi_cell(effective_width, 6, txt=sanitize_text(overview_text))
    pdf.ln(8)
    
    stock_colors = {symbol1: 'blue', symbol2: 'red'}
    metrics_data = {}
    
    for symbol in [symbol1, symbol2]:
        alert = results[symbol]
        chart_data = charts[symbol]
        state = states[symbol]
        
        pdf.set_font("Arial", 'B', size=13)
        pdf.cell(effective_width, 8, txt=sanitize_text(f"Analysis for {symbol}"), ln=True)
        pdf.ln(3)
        
        currency = "Rs." if symbol.endswith(('.NS', '.BO')) else "$"
        alert = alert.replace("₹", currency)
        
        raw_data = state.get('raw_data', {}) or {}
        indicators = raw_data.get("indicators", {})
        
        confidence = state.get('confidence', 50)
        rsi = indicators.get('rsi', 50.0)
        macd = indicators.get('macd', 0.0)
        current_price = indicators.get('current_price', 0.0)
        
        signal = "Yellow: 'Wait and Watch'"
        if confidence >= 70:
            signal = "Green: 'Yes, Go Ahead!'"
        elif confidence >= 40:
            signal = "Yellow: 'Wait and Watch'"
        else:
            signal = "Red: 'No, Stop!'"
        
        metrics_data[symbol] = {"Confidence": confidence, "RSI": rsi, "MACD": macd, "Current Price": current_price, "Signal": signal}
        
        pdf.set_font("Arial", 'B', size=11)
        pdf.multi_cell(effective_width, 6, txt=sanitize_text(f"Current Price: {currency}{current_price:,.2f}"))
        pdf.set_font("Arial", size=11)
        
        pdf.write(6, sanitize_text("Action: "))
        pdf.set_font("Arial", 'B', size=11)
        pdf.write(6, sanitize_text(f"{action.capitalize()} ({horizon_desc})"))
        pdf.set_font("Arial", size=11)
        pdf.ln(6)
        
        pdf.write(6, sanitize_text("Target Date: "))
        pdf.set_font("Arial", 'B', size=11)
        pdf.write(6, sanitize_text(f"{target_date}"))
        pdf.set_font("Arial", size=11)
        pdf.ln(6)
        
        pdf.write(6, sanitize_text("Signal: "))
        pdf.set_font("Arial", 'B', size=11)
        pdf.write(6, sanitize_text(signal))
        pdf.set_font("Arial", size=11)
        pdf.ln(6)
        
        pdf.write(6, sanitize_text("Confidence: "))
        pdf.set_font("Arial", 'B', size=11)
        pdf.write(6, sanitize_text(f"{confidence}%"))
        pdf.set_font("Arial", size=11)
        pdf.ln(8)
        
        pdf.set_font("Arial", 'B', size=12)
        pdf.cell(effective_width, 8, txt=sanitize_text("Key Metrics"), ln=True)
        pdf.ln(2)
        
        col1_width = effective_width * 0.6
        col2_width = effective_width * 0.4
        
        pdf.set_font("Arial", 'B', size=10)
        pdf.set_fill_color(200, 220, 255)
        pdf.cell(col1_width, 8, txt=sanitize_text("Metric"), border=1, align='C', fill=True)
        pdf.cell(col2_width, 8, txt=sanitize_text("Value"), border=1, align='C', fill=True)
        pdf.ln()
        
        pdf.set_font("Arial", size=10)
        for metric in ["Confidence", "RSI", "MACD", "Current Price"]:
            value = metrics_data[symbol][metric]
            pdf.cell(col1_width, 8, txt=sanitize_text(metric), border=1)
            
            if metric == "Confidence":
                pdf.set_font("Arial", 'B', size=10)
                pdf.cell(col2_width, 8, txt=sanitize_text(f"{value}%"), border=1, align='C')
                pdf.set_font("Arial", size=10)
            else:
                pdf.cell(col2_width, 8, txt=sanitize_text(str(value)), border=1, align='C')
            pdf.ln()
        
        pdf.ln(6)
        
        if chart_data:
            dates = [datetime.strptime(point['date'], '%Y-%m-%d') for point in chart_data]
            prices = [point['price'] for point in chart_data]
            
            plt.figure(figsize=(9, 4.5), dpi=100)
            plt.plot(dates, prices, label=f'{symbol} Price', color=stock_colors[symbol], linewidth=2.5, marker='o', markersize=3)
            plt.title(f'{symbol} Price Trend ({horizon_desc})', fontsize=14, fontweight='bold')
            plt.xlabel('Date', fontsize=11, fontweight='bold')
            plt.ylabel('Price', fontsize=11, fontweight='bold')
            plt.legend(fontsize=10)
            plt.grid(True, alpha=0.3, linestyle='--')
            plt.xticks(rotation=45)
            plt.tight_layout()
            
            temp_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
            plt.savefig(temp_img.name, dpi=100, bbox_inches='tight')
            plt.close()
            
            pdf.image(temp_img.name, x=6.35, w=effective_width)
            pdf.ln(8)
        
        pdf.set_font("Arial", size=10)
        if "Should I Buy?" in alert or "Should I Sell?" in alert:
            try:
                search_term = "Should I Buy?" if "Should I Buy?" in alert else "Should I Sell?"
                should_section = alert.split(search_term)[1].split("Best Date/Time")[0].strip()
                pdf.set_font("Arial", 'B', size=11)
                pdf.cell(effective_width, 7, txt=sanitize_text(f"Should I {action.capitalize()}?"), ln=True)
                pdf.set_font("Arial", size=10)
                pdf.multi_cell(effective_width, 6, txt=sanitize_text(should_section[:300]))
                pdf.ln(3)
            except: pass
        
        pdf.set_font("Arial", 'B', size=11)
        pdf.cell(effective_width, 7, txt=sanitize_text("Why Buy/Sell?"), ln=True)
        pdf.set_font("Arial", size=10)
        why_text = f"Based on indicators (RSI {rsi:.1f}, MACD {macd:.2f}), {action} is {'favorable' if confidence > 60 else 'mixed' if confidence > 40 else 'unfavorable'} due to trends and news sentiment."
        pdf.multi_cell(effective_width, 6, txt=sanitize_text(why_text))
        pdf.ln(3)
        
        pdf.set_font("Arial", 'B', size=11)
        pdf.cell(effective_width, 7, txt=sanitize_text("When?"), ln=True)
        pdf.set_font("Arial", size=10)
        when_text = f"Optimal time: Target date {target_date} based on {horizon_desc} strategy for potential gains."
        pdf.multi_cell(effective_width, 6, txt=sanitize_text(when_text))
        pdf.ln(3)
        
        pdf.set_font("Arial", 'B', size=11)
        pdf.cell(effective_width, 7, txt=sanitize_text("Which Trade?"), ln=True)
        pdf.set_font("Arial", size=10)
        trade_type = 'short-term' if horizon in ['intraday', 'scalping'] else 'medium-term' if horizon in ['swing', 'momentum'] else 'long-term'
        which_text = f"Trade Type: {horizon.replace('_', ' ').title()} ({action}). Suitable for {trade_type} investors."
        pdf.multi_cell(effective_width, 6, txt=sanitize_text(which_text))
        pdf.ln(3)
        
        pdf.set_font("Arial", 'B', size=11)
        pdf.cell(effective_width, 7, txt=sanitize_text("Risks?"), ln=True)
        pdf.set_font("Arial", size=10)
        rsi_status = 'overbought' if rsi > 70 else 'oversold' if rsi < 30 else 'neutral'
        risks_text = f"Market volatility, news sentiment changes, potential reversals (e.g., RSI {rsi:.1f} indicates {rsi_status}). Allocate cautiously based on confidence."
        pdf.multi_cell(effective_width, 6, txt=sanitize_text(risks_text))
        pdf.ln(10)
    
    pdf.set_font("Arial", 'B', size=14)
    pdf.cell(effective_width, 10, txt=sanitize_text("Direct Comparison"), ln=True)
    pdf.set_font("Arial", size=11)
    comparison_text = f"Comparing {symbol1} and {symbol2} based on key metrics. Higher confidence and favorable RSI/MACD indicate stronger signals."
    pdf.multi_cell(effective_width, 6, txt=sanitize_text(comparison_text))
    pdf.ln(6)
    
    col_metric = effective_width * 0.25
    col_val1 = effective_width * 0.25
    col_val2 = effective_width * 0.25
    col_better = effective_width * 0.25
    
    pdf.set_font("Arial", 'B', size=10)
    pdf.set_fill_color(200, 220, 255)
    pdf.cell(col_metric, 8, txt=sanitize_text("Metric"), border=1, align='C', fill=True)
    pdf.cell(col_val1, 8, txt=sanitize_text(f"{symbol1}"), border=1, align='C', fill=True)
    pdf.cell(col_val2, 8, txt=sanitize_text(f"{symbol2}"), border=1, align='C', fill=True)
    pdf.cell(col_better, 8, txt=sanitize_text("Better"), border=1, align='C', fill=True)
    pdf.ln()
    
    pdf.set_font("Arial", size=10)
    for metric in ["Confidence", "RSI", "MACD", "Current Price"]:
        val1 = metrics_data[symbol1][metric]
        val2 = metrics_data[symbol2][metric]
        
        if metric == "Confidence":
            better = symbol1 if val1 > val2 else symbol2
        elif metric == "RSI":
            if action == "buy":
                better = symbol1 if (30 <= val1 <= 50) and not (30 <= val2 <= 50) else symbol2 if (30 <= val2 <= 50) else symbol1 if abs(val1 - 40) < abs(val2 - 40) else symbol2
            else:
                better = symbol1 if val1 > 60 else symbol2 if val2 > 60 else symbol1 if val1 > val2 else symbol2
        elif metric == "MACD":
            if action == "buy":
                better = symbol1 if val1 > val2 else symbol2
            else:
                better = symbol1 if val1 < val2 else symbol2
        else:
            better = symbol1 if val1 > val2 else symbol2
        
        pdf.cell(col_metric, 8, txt=sanitize_text(metric), border=1)
        
        if metric == "Confidence":
            pdf.set_font("Arial", 'B', size=10)
            pdf.cell(col_val1, 8, txt=sanitize_text(f"{val1}%"), border=1, align='C')
            pdf.cell(col_val2, 8, txt=sanitize_text(f"{val2}%"), border=1, align='C')
            pdf.set_font("Arial", size=10)
        else:
            pdf.cell(col_val1, 8, txt=sanitize_text(str(val1)), border=1, align='C')
            pdf.cell(col_val2, 8, txt=sanitize_text(str(val2)), border=1, align='C')
        
        pdf.cell(col_better, 8, txt=sanitize_text(better), border=1, align='C')
        pdf.ln()
    
    pdf.ln(8)
    
    metrics = ["Confidence", "RSI", "MACD"]
    val1_list = [metrics_data[symbol1][m] for m in metrics]
    val2_list = [metrics_data[symbol2][m] for m in metrics]
    
    plt.figure(figsize=(9, 4.5), dpi=100)
    x = range(len(metrics))
    plt.bar([i - 0.2 for i in x], val1_list, width=0.4, label=symbol1, color='#4472C4', edgecolor='black')
    plt.bar([i + 0.2 for i in x], val2_list, width=0.4, label=symbol2, color='#ED7D31', edgecolor='black')
    plt.xticks(x, metrics, fontsize=11, fontweight='bold')
    plt.title("Key Metrics Comparison", fontsize=14, fontweight='bold')
    plt.ylabel("Value", fontsize=11, fontweight='bold')
    plt.legend(fontsize=10)
    plt.grid(True, axis='y', alpha=0.3, linestyle='--')
    plt.tight_layout()
    
    temp_bar_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
    plt.savefig(temp_bar_img.name, dpi=100, bbox_inches='tight')
    plt.close()
    
    pdf.image(temp_bar_img.name, x=6.35, w=effective_width)
    pdf.ln(8)
    
    if charts[symbol1] and charts[symbol2]:
        dates1 = [datetime.strptime(point['date'], '%Y-%m-%d') for point in charts[symbol1]]
        prices1 = [point['price'] for point in charts[symbol1]]
        dates2 = [datetime.strptime(point['date'], '%Y-%m-%d') for point in charts[symbol2]]
        prices2 = [point['price'] for point in charts[symbol2]]
        
        plt.figure(figsize=(9, 4.5), dpi=100)
        plt.plot(dates1, prices1, label=f'{symbol1} Price', color='#4472C4', linewidth=2.5, marker='o', markersize=3)
        plt.plot(dates2, prices2, label=f'{symbol2} Price', color='#ED7D31', linewidth=2.5, marker='s', markersize=3)
        plt.title(f'Combined Price Trend Comparison ({horizon_desc})', fontsize=14, fontweight='bold')
        plt.xlabel('Date', fontsize=11, fontweight='bold')
        plt.ylabel('Price', fontsize=11, fontweight='bold')
        plt.legend(fontsize=10)
        plt.grid(True, alpha=0.3, linestyle='--')
        plt.xticks(rotation=45)
        plt.tight_layout()
        
        temp_combined_img = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        plt.savefig(temp_combined_img.name, dpi=100, bbox_inches='tight')
        plt.close()
        
        pdf.image(temp_combined_img.name, x=6.35, w=effective_width)
        pdf.ln(8)
    
    pdf.set_font("Arial", 'B', size=13)
    pdf.cell(effective_width, 8, txt=sanitize_text("Recommendation"), ln=True)
    pdf.ln(2)
    
    better_stock = symbol1 if metrics_data[symbol1]["Confidence"] > metrics_data[symbol2]["Confidence"] else symbol2
    better_confidence = metrics_data[better_stock]["Confidence"]
    
    pdf.set_font("Arial", size=11)
    pdf.write(6, sanitize_text(f"Based on higher confidence ("))
    pdf.set_font("Arial", 'B', size=11)
    pdf.write(6, sanitize_text(f"{better_confidence}%"))
    pdf.set_font("Arial", size=11)
    pdf.write(6, sanitize_text(f") and metrics, "))
    pdf.set_font("Arial", 'B', size=11)
    pdf.write(6, sanitize_text(f"{better_stock}"))
    pdf.set_font("Arial", size=11)
    pdf.write(6, sanitize_text(f" is recommended for {action}. "))
    pdf.set_font("Arial", 'I', size=10)
    pdf.write(6, sanitize_text("Not financial advice-consult a professional."))
    
    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".pdf")
    pdf.output(temp_file.name)
    return temp_file.name, results, charts



# Initialize session state FIRST
if 'query_processed' not in st.session_state:
    st.session_state.query_processed = False
    st.session_state.is_comparison = False
    st.session_state.pdf_path = None
    st.session_state.results = None
    st.session_state.selected_example = ""

# Header
def load_image_base64(path):
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


BASE_DIR = os.path.dirname(os.path.abspath(__file__))
logo_path = os.path.join(BASE_DIR, "static", "main-logo.png")
logo_b64 = load_image_base64(logo_path)

st.markdown(f"""
<div class="header" style="display:flex;gap:64px;align-items:center;justify-content:center">
    <img src="data:image/jpeg;base64,{logo_b64}" width="130" style='display:flex; justify-content:start'/>
    <h1>SparkTradeAnalyzer - Intelligent Stock Analysis</h1>
</div>
""", unsafe_allow_html=True)

# Subtitle
st.markdown("""
<div style="text-align: center; padding: 20px 0; color: rgba(255,255,255,0.8);">
    <p style="font-size: 18px; margin: 0;">Smart Trading Decisions with Live Market Insights</p>
    <p style="font-size: 14px; margin: 5px 0 0 0; opacity: 0.7;">Powered by SparkBrains</p>
</div>
""", unsafe_allow_html=True)

# How to Use Section
st.markdown('<div class="how-to-container">', unsafe_allow_html=True)
st.markdown("### How to Use")
st.markdown("""
<ul style="display: flex; flex-wrap: wrap; gap: 0px; list-style: none; padding: 0; margin: 0; color: white;">
    <li>○ Analyze stock between markets</li>
    <li>○ Add trading details</li>
    <li>○ Optimize multi-stock portfolios</li>
    <li>○ Check market or conditions</li>
    <li>○ Save trading preferences</li>
</ul>
""", unsafe_allow_html=True)
st.markdown('</div>', unsafe_allow_html=True)

# Main Query Container
st.markdown('<div class="query-container">', unsafe_allow_html=True)
col1, col2 = st.columns(2)

with col1:
    st.markdown("### Query Input")
    
    # Use selected example if available, otherwise empty
    default_value = st.session_state.selected_example if st.session_state.selected_example else ""
    
    user_query = st.text_area(
        "",
        value=default_value,
        placeholder="Examples:\n• Plan route from Mumbai to Pune\n• Route from Delhi to Jaipur with 500kg package\n• What was my last route?\n• Save preference: I prefer morning deliveries",
        height=200,
        label_visibility="collapsed",
        key="query_input"
    )
    
    col_btn1, col_btn2 = st.columns(2)
    with col_btn1:
        send_btn = st.button("Send", use_container_width=True, type="primary")
    with col_btn2:
        clear_btn = st.button("Clear", use_container_width=True, type="secondary")
    
    st.markdown('</div>', unsafe_allow_html=True)

with col2:
    st.markdown("### Quick Examples")
    
    examples = [
        "Should I sell my intraday position in TMPV.NS today?",
        "For intraday trading, should I buy IDBI Bank or SBI Bank today?",
        "I want to buy Infosys this week and sell it after 2 weeks.",
        "Is Google stock good for a 10-day swing trade?",
        "Should I buy ICICI Bank now to sell in the middle of next month?",
        "I purchased TCS stock 3 months ago, should I sell it today?",
        "I want to enter a momentum trade on Bharti Airtel for the next month."
    ]
    
    for idx, example in enumerate(examples):
        if st.button(example, key=f"example_{idx}", use_container_width=True):
            st.session_state.selected_example = example
            st.rerun()
    
    st.markdown('</div>', unsafe_allow_html=True)

st.markdown('</div>', unsafe_allow_html=True)

# Handle clear button
if clear_btn:
    st.session_state.query_processed = False
    st.session_state.is_comparison = False
    st.session_state.pdf_path = None
    st.session_state.results = None
    st.session_state.selected_example = ""
    st.rerun()

# Handle send button
if send_btn and user_query:
    # Clear the selected example after using it
    if st.session_state.selected_example:
        st.session_state.selected_example = ""
    
    with st.spinner("Analyzing..."):
        parsed = classify_and_parse_query(user_query)
        
        if parsed["is_comparison"] and len(parsed["symbols"]) == 2:
            st.session_state.is_comparison = True
            symbols = parsed["symbols"]
            details = parsed.get("other_details", {})
            action = details.get("action", "buy")
            horizon = details.get("horizon", "intraday")
            date = details.get("date")
            holding_period = details.get("holding_period")
            
            with st.spinner(f"Comparing {symbols[0]} vs {symbols[1]}..."):
                pdf_path, results, charts = process_comparison_query(symbols[0], symbols[1], action, horizon, date, holding_period)
                st.session_state.pdf_path = pdf_path
                st.session_state.results = results
                st.session_state.query_processed = True
                st.success("✅ Comparison complete!")
        
        elif not parsed["is_comparison"] and len(parsed["symbols"]) == 1:
            st.session_state.is_comparison = False
            symbol = parsed["symbols"][0]
            details = parsed.get("other_details", {})
            action = details.get("action", "buy")
            horizon = details.get("horizon", "intraday")
            date = details.get("date")
            holding_period = details.get("holding_period")
            
            with st.spinner(f"Analyzing {symbol}..."):
                alert, chart_data, final_state = process_normal_query(symbol, action, horizon, date, holding_period)
                
                st.markdown('<div class="results-heading block">', unsafe_allow_html=True)
                st.markdown("## Analysis Report")
                st.markdown('</div>', unsafe_allow_html=True)
                
                st.text_area("Report", alert, height=400, label_visibility="collapsed")
                
                if chart_data:
                    st.markdown("### Price Trend")
                    dates = [datetime.strptime(p['date'], '%Y-%m-%d') for p in chart_data]
                    prices = [p['price'] for p in chart_data]
                    
                    fig, ax = plt.subplots(figsize=(10, 5))
                    ax.plot(dates, prices, label=f'{symbol}', linewidth=2.5, marker='o', markersize=4, color='#78C841')
                    ax.set_title(f'{symbol} Price Trend', fontsize=14, fontweight='bold')
                    ax.set_xlabel('Date')
                    ax.set_ylabel('Price')
                    ax.legend()
                    ax.grid(alpha=0.3)
                    plt.xticks(rotation=45)
                    plt.tight_layout()
                    st.pyplot(fig)
                
                st.session_state.query_processed = True
        else:
            st.error("❌ Could not parse query. Please try again.")

# Display download button if comparison is complete
if st.session_state.query_processed and st.session_state.is_comparison and st.session_state.pdf_path:
    st.markdown('<div class="results-heading block">', unsafe_allow_html=True)
    st.markdown("## Results")
    st.markdown('</div>', unsafe_allow_html=True)
    
    with open(st.session_state.pdf_path, "rb") as pdf_file:
        st.download_button(
            "📥 Download Comparison Report (PDF)",
            pdf_file,
            "stock_comparison.pdf",
            "application/pdf",
            use_container_width=True
        )