import os
import sys
import subprocess
import re
import multiprocessing
import atexit
from collections.abc import Iterator
from functools import lru_cache
import datetime
import time
import gradio as gr
import gradio.themes as themes
from huggingface_hub import hf_hub_download, login
import logging
import pandas as pd
import torch

# Set up logging with more detail
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Install llama-cpp-python with appropriate backend
try:
    from llama_cpp import Llama
except ModuleNotFoundError:
    if torch.cuda.is_available():
        logger.info("Installing llama-cpp-python with CUDA support.")
        os.environ['CMAKE_ARGS'] = "-DLLAMA_CUDA=ON"
        subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python", "--force-reinstall", "--upgrade", "--no-cache-dir"])
    else:
        logger.info("Installing llama-cpp-python without additional flags.")
        subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python", "--force-reinstall", "--upgrade", "--no-cache-dir"])
from llama_cpp import Llama

# Install yfinance if not present
try:
    import yfinance as yf
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "yfinance"])
import yfinance as yf

# Additional imports for visualization
try:
    import matplotlib.pyplot as plt
    from PIL import Image
    import io
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "matplotlib", "pillow"])
import matplotlib.pyplot as plt
from PIL import Image
import io

# Constants
MAX_MAX_NEW_TOKENS = 512
DEFAULT_MAX_NEW_TOKENS = 256
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "1024"))
YFINANCE_TIMEOUT = 10
CACHE_EXPIRY_HOURS = 1

DESCRIPTION = """# FinChat: Investing Q&A (Optimized for Speed)
This application delivers an interactive chat interface powered by a highly efficient, small AI model adapted for addressing investing and finance inquiries through specialized prompt engineering. It ensures rapid, reasoned responses to user queries.
<p>Running on CPU or GPU if available. Using Phi-2 model for faster inference. Includes performance optimizations with caching and improved error handling.</p>"""

LICENSE = """<p/>
---
This application employs the Phi-2 model, governed by Microsoft's Terms of Use. Refer to the [model card](https://huggingface.co/TheBloke/phi-2-GGUF) for details."""

DEFAULT_SYSTEM_PROMPT = """You are FinChat, a knowledgeable AI assistant specializing in investing and finance. Provide accurate, helpful, reasoned, and concise answers to investing questions. Always base responses on reliable information and advise users to consult professionals for personalized advice.
Always respond exclusively in English. Use bullet points for clarity.
Example:
User: average return for TSLA between 2010 and 2020
Assistant:
- TSLA CAGR (2010-2020): ~63.01%
- Represents average annual return with compounding
- Past performance not indicative of future results
- Consult a financial advisor"""

# Company name to ticker mapping
COMPANY_TO_TICKER = {
    "opendoor": "OPEN",
    "tesla": "TSLA",
    "apple": "AAPL",
    "amazon": "AMZN",
    "microsoft": "MSFT",
    "google": "GOOGL",
    "facebook": "META",
    "meta": "META",
    "nvidia": "NVDA",
    "netflix": "NFLX",
}

# Compiled regex patterns for better performance
CAGR_PATTERN = re.compile(
    r'(?:average\s+return|cagr)\s+(?:for\s+)?([\w\s,]+(?:and\s+[\w\s,]+)?)\s+(?:between|from)\s+(\d{4})\s+(?:and|to)\s+(\d{4})',
    re.IGNORECASE
)

COMPOUND_INTEREST_PATTERN = re.compile(
    r'(?:save|invest|deposit)\s*\$?([\d,]+(?:\.\d+)?)\s*(?:right now|today)?\s*(?:under|at)\s*([\d.]+)%\s*(?:interest|rate)?\s*(?:annually|per year)?\s*(?:over|for)\s*(\d+)\s*years?',
    re.IGNORECASE
)

# Load the model
try:
    model_path = hf_hub_download(
        repo_id="TheBloke/phi-2-GGUF",
        filename="phi-2.Q4_K_M.gguf"
    )
    n_gpu_layers = -1 if torch.cuda.is_available() else 0
    llm = Llama(
        model_path=model_path,
        n_ctx=1024,
        n_batch=1024,
        n_threads=multiprocessing.cpu_count(),
        n_gpu_layers=n_gpu_layers,
        chat_format="chatml"
    )
    logger.info(f"Model loaded successfully with n_gpu_layers={n_gpu_layers}.")
    
    # Warm up the model
    llm("Warm-up prompt", max_tokens=1, echo=False)
    logger.info("Model warm-up completed.")
except Exception as e:
    logger.error(f"Error loading model: {str(e)}")
    raise

atexit.register(llm.close)

# Cache for stock data with timestamp
_stock_cache = {}
_cache_timestamps = {}

def sanitize_ticker(ticker):
    """Sanitize ticker input to prevent injection and validate format"""
    if not ticker or not isinstance(ticker, str):
        return None
    cleaned = re.sub(r'[^A-Z0-9.-]', '', ticker.upper().strip())
    if len(cleaned) > 10 or len(cleaned) == 0:
        return None
    return cleaned

def validate_year_range(start_year, end_year):
    """Validate year inputs"""
    try:
        start = int(start_year)
        end = int(end_year)
        current_year = datetime.datetime.now().year
        
        if not (1900 <= start <= current_year):
            return False, f"Start year must be between 1900 and {current_year}"
        if not (1900 <= end <= current_year + 1):
            return False, f"End year must be between 1900 and {current_year + 1}"
        if end <= start:
            return False, "End year must be after start year"
        
        return True, "Valid"
    except ValueError:
        return False, "Years must be valid integers"

@lru_cache(maxsize=100)
def get_stock_data_cached(ticker, start_date, end_date, cache_key):
    """Cache stock data to avoid repeated API calls. cache_key includes hour for expiry."""
    try:
        logger.info(f"Fetching data for {ticker} from {start_date} to {end_date}")
        data = yf.download(
            ticker, 
            start=start_date, 
            end=end_date, 
            progress=False, 
            auto_adjust=False,
            timeout=YFINANCE_TIMEOUT
        )
        return data
    except Exception as e:
        logger.error(f"Error fetching data for {ticker}: {str(e)}")
        return None

def get_current_cache_key():
    """Generate cache key that expires every hour"""
    now = datetime.datetime.now()
    return f"{now.year}{now.month}{now.day}{now.hour}"

def calculate_cagr(ticker, start_year, end_year):
    """Calculate CAGR for a ticker with error handling"""
    try:
        # Sanitize ticker
        clean_ticker = sanitize_ticker(ticker)
        if not clean_ticker:
            return f"- {ticker}: Invalid ticker format"
        
        # Validate years
        valid, msg = validate_year_range(start_year, end_year)
        if not valid:
            return f"- {clean_ticker}: {msg}"
        
        # Get cached data
        cache_key = get_current_cache_key()
        data = get_stock_data_cached(
            clean_ticker,
            f"{start_year}-01-01",
            f"{end_year}-12-31",
            cache_key
        )
        
        if data is None:
            return f"- {clean_ticker}: Error fetching data (API timeout or network issue)"
        
        if data.empty:
            return f"- {clean_ticker}: No historical data available between {start_year} and {end_year}"
        
        # Handle MultiIndex columns
        if isinstance(data.columns, pd.MultiIndex):
            data.columns = data.columns.droplevel(1)
        
        # Check for Adj Close column
        if 'Adj Close' not in data.columns:
            return f"- {clean_ticker}: Adjusted Close price data not available"
        
        # Calculate CAGR
        initial = data['Adj Close'].iloc[0]
        final = data['Adj Close'].iloc[-1]
        
        if pd.isna(initial) or pd.isna(final):
            return f"- {clean_ticker}: Missing price data"
        
        if initial <= 0 or final <= 0:
            return f"- {clean_ticker}: Invalid price data (negative or zero values)"
        
        start_date = data.index[0]
        end_date = data.index[-1]
        days = (end_date - start_date).days
        years = days / 365.25
        
        if years <= 0:
            return f"- {clean_ticker}: Invalid time period"
        
        cagr = ((final / initial) ** (1 / years) - 1) * 100
        
        # Add context about data quality
        actual_start = start_date.strftime('%Y-%m-%d')
        actual_end = end_date.strftime('%Y-%m-%d')
        date_note = f" (data: {actual_start} to {actual_end})" if actual_start != f"{start_year}-01-01" else ""
        
        return f"- {clean_ticker}: ~{cagr:.2f}%{date_note}"
        
    except Exception as e:
        logger.error(f"Unexpected error calculating CAGR for {ticker}: {str(e)}")
        return f"- {ticker}: Calculation error - {str(e)}"

def generate(
    message: str,
    chat_history: list[dict],
    system_prompt: str = DEFAULT_SYSTEM_PROMPT,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    start_time = time.time()
    logger.info(f"Generating response for message: {message[:100]}...")
    
    lower_message = message.lower().strip()
    
    # Quick responses for common queries
    if lower_message in ["hi", "hello", "hey"]:
        response = "Hello! I'm FinChat, your financial advisor. Ask me anything about investing, stocks, CAGR, compound interest, or portfolio analysis!"
        logger.info(f"Quick response generated in {time.time() - start_time:.2f}s")
        yield response
        return
    
    if "what is cagr" in lower_message:
        response = """- CAGR stands for Compound Annual Growth Rate.
- It measures the mean annual growth rate of an investment over a specified period longer than one year, accounting for compounding.
- Formula: CAGR = (Ending Value / Beginning Value)^(1 / Number of Years) - 1
- Useful for comparing investments over time.
- Past performance is not indicative of future results. Consult a financial advisor."""
        logger.info(f"Quick response generated in {time.time() - start_time:.2f}s")
        yield response
        return
    
    # Compound interest calculation
    compound_match = COMPOUND_INTEREST_PATTERN.search(lower_message)
    if compound_match:
        try:
            principal_str = compound_match.group(1).replace(',', '')
            principal = float(principal_str)
            rate = float(compound_match.group(2)) / 100
            years = int(compound_match.group(3))
            
            if principal <= 0:
                yield "Error: Principal amount must be positive."
                return
            if rate < 0 or rate > 1:
                yield "Error: Interest rate must be between 0% and 100%."
                return
            if years <= 0 or years > 100:
                yield "Error: Years must be between 1 and 100."
                return
            
            balance = principal * (1 + rate) ** years
            total_interest = balance - principal
            
            response = (
                f"**Compound Interest Calculation**\n\n"
                f"- Starting Principal: ${principal:,.2f}\n"
                f"- Annual Interest Rate: {rate*100:.2f}%\n"
                f"- Time Period: {years} years\n"
                f"- Compounding: Annually\n\n"
                f"**Results:**\n"
                f"- Final Balance (Year {years}): ${balance:,.2f}\n"
                f"- Total Interest Earned: ${total_interest:,.2f}\n"
                f"- Total Growth: {((balance/principal - 1) * 100):.2f}%\n\n"
                f"*Note: This assumes annual compounding with no additional deposits or withdrawals. Actual results may vary. Consult a financial advisor.*"
            )
            logger.info(f"Compound interest calculated in {time.time() - start_time:.2f}s")
            yield response
            return
        except ValueError as ve:
            logger.error(f"Error parsing compound interest query: {str(ve)}")
            yield "Error: Please ensure amount, rate, and years are valid numbers. Example: 'If I save $10000 at 5% interest over 10 years'"
            return
        except Exception as e:
            logger.error(f"Unexpected error in compound interest: {str(e)}")
            yield f"Error calculating compound interest: {str(e)}"
            return
    
    # CAGR calculation with improved pattern matching
    cagr_match = CAGR_PATTERN.search(lower_message)
    if cagr_match:
        tickers_str, start_year, end_year = cagr_match.groups()
        tickers = [t.strip() for t in re.split(r',|\band\b', tickers_str) if t.strip()]
        
        # Apply company-to-ticker mapping
        mapped_tickers = []
        for ticker in tickers:
            lower_ticker = ticker.lower()
            if lower_ticker in COMPANY_TO_TICKER:
                mapped_tickers.append(COMPANY_TO_TICKER[lower_ticker])
            else:
                mapped_tickers.append(ticker.upper())
        
        # Validate year range first
        valid, msg = validate_year_range(start_year, end_year)
        if not valid:
            yield f"Error: {msg}"
            return
        
        if len(mapped_tickers) > 10:
            yield "Error: Too many tickers requested. Please limit to 10 tickers per query."
            return
        
        responses = []
        for ticker in mapped_tickers:
            result = calculate_cagr(ticker, start_year, end_year)
            responses.append(result)
        
        full_response = (
            f"**CAGR Analysis ({start_year} - {end_year})**\n\n"
            + "\n".join(responses) + 
            "\n\n*Notes:*\n"
            "- CAGR represents average annual returns with compounding\n"
            "- Based on adjusted closing prices\n"
            "- Past performance is not indicative of future results\n"
            "- Please consult a financial advisor for investment decisions"
        )
        
        logger.info(f"CAGR response generated in {time.time() - start_time:.2f}s")
        yield full_response
        return
    
    # Build conversation for LLM
    conversation = [{"role": "system", "content": system_prompt}]
    for msg in chat_history[-3:]:
        if msg["role"] in ["user", "assistant"]:
            conversation.append({"role": msg["role"], "content": msg["content"]})
    conversation.append({"role": "user", "content": message})
    
    # Token length check with truncation
    prompt_text = "\n".join(d["content"] for d in conversation)
    input_tokens = llm.tokenize(prompt_text.encode("utf-8"), add_bos=False)
    
    while len(input_tokens) > MAX_INPUT_TOKEN_LENGTH:
        logger.warning(f"Input tokens ({len(input_tokens)}) exceed limit. Truncating history.")
        if len(conversation) > 2:
            conversation.pop(1)
            prompt_text = "\n".join(d["content"] for d in conversation)
            input_tokens = llm.tokenize(prompt_text.encode("utf-8"), add_bos=False)
        else:
            yield "Error: Input is too long. Please shorten your query or start a new conversation."
            return
    
    # Generate response
    try:
        response = ""
        sentence_buffer = ""
        stream = llm.create_chat_completion(
            messages=conversation,
            max_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
            repeat_penalty=repetition_penalty,
            stream=True
        )
        
        sentence_endings = ['.', '!', '?']
        
        for chunk in stream:
            delta = chunk["choices"][0]["delta"]
            if "content" in delta and delta["content"] is not None:
                cleaned_chunk = re.sub(
                    r'<\|(?:im_start|im_end|system|user|assistant)\|>|</s>|\[END\]', 
                    '', 
                    delta["content"]
                )
                if not cleaned_chunk:
                    continue
                
                sentence_buffer += cleaned_chunk
                response += cleaned_chunk
                
                if any(sentence_buffer.strip().endswith(ending) for ending in sentence_endings):
                    yield response
                    sentence_buffer = ""
            
            if chunk["choices"][0]["finish_reason"] is not None:
                if sentence_buffer.strip():
                    yield response
                break
        
        duration = time.time() - start_time
        logger.info(f"LLM response generated in {duration:.2f}s")
        
    except ValueError as ve:
        if "exceed context window" in str(ve):
            yield "Error: Input exceeds context window. Please try a shorter query."
        else:
            logger.error(f"ValueError during generation: {str(ve)}")
            yield f"Error: {str(ve)}"
    except Exception as e:
        logger.error(f"Error during response generation: {str(e)}")
        yield f"Error generating response. Please try again or rephrase your question."

def process_portfolio(df, growth_rate):
    """Process portfolio with enhanced error handling and validation"""
    if df is None or len(df) == 0:
        return "", None
    
    try:
        if not isinstance(df, pd.DataFrame):
            df = pd.DataFrame(df, columns=["Ticker", "Shares", "Avg Cost", "Current Price"])
        
        # Validate and convert numeric columns
        for col in ["Shares", "Avg Cost", "Current Price"]:
            df[col] = pd.to_numeric(df[col], errors='coerce')
        
        df = df.dropna(subset=["Ticker"])
        
        portfolio = {}
        errors = []
        
        for idx, row in df.iterrows():
            ticker = sanitize_ticker(row["Ticker"]) if pd.notna(row["Ticker"]) else None
            if not ticker:
                continue
            
            shares = float(row["Shares"]) if pd.notna(row["Shares"]) else 0
            cost = float(row["Avg Cost"]) if pd.notna(row["Avg Cost"]) else 0
            price = float(row["Current Price"]) if pd.notna(row["Current Price"]) else 0
            
            if shares <= 0:
                errors.append(f"{ticker}: Invalid shares count")
                continue
            if price < 0 or cost < 0:
                errors.append(f"{ticker}: Negative prices not allowed")
                continue
            
            value = shares * price
            cost_basis = shares * cost
            gain_loss = value - cost_basis
            gain_loss_pct = (gain_loss / cost_basis * 100) if cost_basis > 0 else 0
            
            portfolio[ticker] = {
                'shares': shares,
                'cost': cost,
                'price': price,
                'value': value,
                'cost_basis': cost_basis,
                'gain_loss': gain_loss,
                'gain_loss_pct': gain_loss_pct
            }
        
        if not portfolio:
            return "No valid portfolio entries found. Please check your data.", None
        
        total_value_now = sum(v['value'] for v in portfolio.values())
        total_cost_basis = sum(v['cost_basis'] for v in portfolio.values())
        total_gain_loss = total_value_now - total_cost_basis
        total_gain_loss_pct = (total_gain_loss / total_cost_basis * 100) if total_cost_basis > 0 else 0
        
        allocations = {k: v['value'] / total_value_now for k, v in portfolio.items()} if total_value_now > 0 else {}
        
        # Create allocation pie chart
        fig_alloc, ax_alloc = plt.subplots(figsize=(8, 6))
        colors = plt.cm.Set3(range(len(allocations)))
        ax_alloc.pie(
            allocations.values(), 
            labels=allocations.keys(), 
            autopct='%1.1f%%',
            colors=colors,
            startangle=90
        )
        ax_alloc.set_title('Portfolio Allocation by Value', fontsize=14, fontweight='bold')
        
        buf_alloc = io.BytesIO()
        fig_alloc.savefig(buf_alloc, format='png', bbox_inches='tight', dpi=100)
        buf_alloc.seek(0)
        chart_alloc = Image.open(buf_alloc)
        plt.close(fig_alloc)
        
        # Project future values
        def project_value(value, years, rate):
            return value * (1 + rate / 100) ** years
        
        projections = {
            '1 year': sum(project_value(v['value'], 1, growth_rate) for v in portfolio.values()),
            '2 years': sum(project_value(v['value'], 2, growth_rate) for v in portfolio.values()),
            '5 years': sum(project_value(v['value'], 5, growth_rate) for v in portfolio.values()),
            '10 years': sum(project_value(v['value'], 10, growth_rate) for v in portfolio.values())
        }
        
        # Build detailed report
        data_str = "**📊 Portfolio Analysis**\n\n"
        data_str += "**Current Holdings:**\n"
        for ticker, data in portfolio.items():
            data_str += (
                f"- {ticker}: {data['shares']:.2f} shares @ ${data['price']:.2f} "
                f"(Cost: ${data['cost']:.2f}) = ${data['value']:,.2f} "
                f"[{data['gain_loss_pct']:+.2f}%]\n"
            )
        
        data_str += f"\n**Portfolio Summary:**\n"
        data_str += f"- Total Value: ${total_value_now:,.2f}\n"
        data_str += f"- Total Cost Basis: ${total_cost_basis:,.2f}\n"
        data_str += f"- Total Gain/Loss: ${total_gain_loss:+,.2f} ({total_gain_loss_pct:+.2f}%)\n"
        
        data_str += f"\n**Projected Values (at {growth_rate}% annual growth):**\n"
        for period, value in projections.items():
            gain = value - total_value_now
            data_str += f"- {period}: ${value:,.2f} (+${gain:,.2f})\n"
        
        if errors:
            data_str += f"\n**⚠️ Warnings:**\n"
            for error in errors:
                data_str += f"- {error}\n"
        
        data_str += "\n*Note: Projections assume constant growth rate and no additional contributions. Actual results will vary. Consult a financial advisor.*"
        
        return data_str, chart_alloc
        
    except Exception as e:
        logger.error(f"Error processing portfolio: {str(e)}")
        return f"Error processing portfolio: {str(e)}", None

def fetch_current_prices(df):
    """Fetch current prices with timeout and error handling"""
    if df is None or len(df) == 0:
        return df
    
    try:
        if not isinstance(df, pd.DataFrame):
            df = pd.DataFrame(df, columns=["Ticker", "Shares", "Avg Cost", "Current Price"])
        
        updated_count = 0
        failed_tickers = []
        
        for i in df.index:
            ticker = df.at[i, "Ticker"]
            if pd.notna(ticker) and ticker.strip():
                clean_ticker = sanitize_ticker(ticker)
                if not clean_ticker:
                    failed_tickers.append(f"{ticker} (invalid format)")
                    continue
                
                try:
                    ticker_obj = yf.Ticker(clean_ticker)
                    info = ticker_obj.info
                    price = info.get('currentPrice') or info.get('regularMarketPrice')
                    
                    if price is not None and price > 0:
                        df.at[i, "Current Price"] = price
                        updated_count += 1
                    else:
                        failed_tickers.append(f"{clean_ticker} (no price data)")
                except Exception as e:
                    logger.warning(f"Failed to fetch price for {clean_ticker}: {str(e)}")
                    failed_tickers.append(f"{clean_ticker} ({str(e)[:30]})")
        
        if updated_count > 0:
            logger.info(f"Successfully updated {updated_count} prices")
        if failed_tickers:
            logger.warning(f"Failed to fetch: {', '.join(failed_tickers)}")
        
        return df
        
    except Exception as e:
        logger.error(f"Error in fetch_current_prices: {str(e)}")
        return df

# Gradio interface
with gr.Blocks(theme=themes.Soft(), css="""
    #chatbot {height: 800px; overflow: auto;}
    .performance-note {color: #666; font-size: 0.9em; font-style: italic;}
""") as demo:
    gr.Markdown(DESCRIPTION)
    
    chatbot = gr.Chatbot(label="FinChat", type="messages")
    msg = gr.Textbox(
        label="Ask a finance question", 
        placeholder="e.g., 'What is CAGR?' or 'Average return for AAPL between 2010 and 2020'",
        info="Enter your query. Responses are cached for better performance."
    )
    
    with gr.Row():
        submit = gr.Button("Submit", variant="primary")
        clear = gr.Button("Clear")
    
    gr.Examples(
        examples=[
            "What is CAGR?",
            "Average return for AAPL between 2015 and 2023",
            "Average return for TSLA and NVDA between 2018 and 2023",
            "If I save $10000 at 5% interest over 10 years",
            "Explain compound interest"
        ],
        inputs=msg,
        label="Example Queries"
    )
    
    with gr.Accordion("📈 Enter Portfolio for Projections", open=False):
        portfolio_df = gr.Dataframe(
            headers=["Ticker", "Shares", "Avg Cost", "Current Price"],
            datatype=["str", "number", "number", "number"],
            row_count=5,
            col_count=(4, "fixed"),
            label="Portfolio Data",
            interactive=True
        )
        gr.Markdown("""
        **Instructions:** 
        - Enter stock tickers (e.g., AAPL, TSLA)
        - Fill in number of shares and your average cost per share
        - Click 'Fetch Current Prices' to auto-populate current prices
        - Adjust growth rate for future projections
        """)
        
        fetch_button = gr.Button("🔄 Fetch Current Prices", variant="secondary")
        fetch_button.click(fetch_current_prices, inputs=portfolio_df, outputs=portfolio_df)
        
        growth_rate = gr.Slider(
            minimum=0, 
            maximum=50, 
            step=1, 
            value=10, 
            label="Annual Growth Rate (%)", 
            interactive=True,
            info="Expected annual return for projections (0-50%)"
        )
        growth_rate_label = gr.Markdown("**Selected Growth Rate: 10%**")
    
    with gr.Accordion("⚙️ Advanced Settings", open=False):
        system_prompt = gr.Textbox(
            label="System Prompt", 
            value=DEFAULT_SYSTEM_PROMPT, 
            lines=6,
            info="Customize the AI's behavior"
        )
        temperature = gr.Slider(
            label="Temperature", 
            value=0.6, 
            minimum=0.0, 
            maximum=1.0, 
            step=0.05,
            info="Lower = more focused, Higher = more creative"
        )
        top_p = gr.Slider(
            label="Top P", 
            value=0.9, 
            minimum=0.0, 
            maximum=1.0, 
            step=0.05,
            info="Nucleus sampling threshold"
        )
        top_k = gr.Slider(
            label="Top K", 
            value=50, 
            minimum=1, 
            maximum=100, 
            step=1,
            info="Limit to top K tokens"
        )
        repetition_penalty = gr.Slider(
            label="Repetition Penalty", 
            value=1.2, 
            minimum=1.0, 
            maximum=2.0,
            step=0.05,
            info="Penalize repeated tokens"
        )
        max_new_tokens = gr.Slider(
            label="Max New Tokens", 
            value=DEFAULT_MAX_NEW_TOKENS, 
            minimum=1, 
            maximum=MAX_MAX_NEW_TOKENS, 
            step=1,
            info="Maximum length of generated response"
        )
    
    gr.Markdown(LICENSE)
    gr.Markdown('<p class="performance-note">⚡ Performance optimized with caching and improved error handling</p>')
    
    def update_growth_rate_label(growth_rate):
        return f"**Selected Growth Rate: {growth_rate}%**"
    
    def user(message, history):
        if not message:
            return "", history
        return "", history + [{"role": "user", "content": message}]
    
    def bot(history, sys_prompt, temp, tp, tk, rp, mnt, portfolio_df, growth_rate):
        if not history:
            logger.warning("History is empty, initializing with user message.")
            history = [{"role": "user", "content": ""}]
        
        message = history[-1]["content"]
        portfolio_data, chart_alloc = process_portfolio(portfolio_df, growth_rate)
        
        if portfolio_data:
            message += "\n\n" + portfolio_data
        
        history[-1]["content"] = message
        history.append({"role": "assistant", "content": ""})
        
        for new_text in generate(message, history[:-1], sys_prompt, mnt, temp, tp, tk, rp):
            history[-1]["content"] = new_text
            yield history, f"**Selected Growth Rate: {growth_rate}%**"
        
        if chart_alloc:
            # Append chart as a separate message
            yield history, f"**Selected Growth Rate: {growth_rate}%**"
    
    growth_rate.change(update_growth_rate_label, inputs=growth_rate, outputs=growth_rate_label)
    
    submit.click(
        user, 
        [msg, chatbot], 
        [msg, chatbot], 
        queue=False
    ).then(
        bot, 
        [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens, portfolio_df, growth_rate], 
        [chatbot, growth_rate_label]
    )
    
    msg.submit(
        user, 
        [msg, chatbot], 
        [msg, chatbot], 
        queue=False
    ).then(
        bot, 
        [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens, portfolio_df, growth_rate], 
        [chatbot, growth_rate_label]
    )
    
    clear.click(lambda: [], None, chatbot, queue=False)

demo.queue(max_size=128).launch()