app.py

import gradio as gr
from datetime import datetime
import os
from dotenv import load_dotenv
import requests # Import requests for HTTP calls to Modal backend
import plotly.io as pio # Import plotly.io to deserialize JSON strings to Plotly figures
# from gradio.themes.utils import fonts

# from backend.stock_analyzer import StockAnalyzer
# from backend.config import AppConfig

# Load environment variables
load_dotenv()


# class BuffetBotTheme(gr.themes.Soft):
#     def __init__(self, **kwargs):
#         super().__init__(
#             font=(
#                 fonts.GoogleFont("Quicksand"),
#                 "ui-sans-serif",
#                 "sans-serif",
#             ),
#             **kwargs
#         )


class BuffetBotApp:
    def __init__(self):
        # self.config = AppConfig()
        # self.stock_analyzer = StockAnalyzer()
        self.mcp_server_url = os.getenv("MODAL_MCP_SERVER_URL") # Get Modal MCP server URL from environment variable
        if not self.mcp_server_url:
            raise ValueError("MODAL_MCP_SERVER_URL environment variable not set.")
        
    def create_ui(self):
        """Create the Gradio interface."""
        custom_css = """
        /* General Body and Font */
        body {
            font-family: 'Roboto', 'Helvetica', 'Arial', sans-serif;
            color: #333;
            background-color: #f5f5f5;
        }

        /* Main Container Styling */
        .gradio-container {
            max-width: 1000px; /* Slightly wider container */
            margin: auto;
            padding: 30px;
            box-shadow: 0 4px 15px rgba(0, 0, 0, 0.1);
            border-radius: 12px;
            background-color: #ffffff;
            margin-top: 30px;
            margin-bottom: 30px;
        }

        /* Headings */
        h1 {
            font-size: 2.8em;
            color: #1a237e; /* Darker blue for prominence */
            text-align: center;
            margin-bottom: 25px;
            font-weight: 700; /* Bold */
            letter-spacing: -0.5px;
        }

        h2 {
            font-size: 2.0em;
            color: #3f51b5; /* Medium blue for subheadings */
            border-bottom: 2px solid #e8eaf6; /* Light separator */
            padding-bottom: 10px;
            margin-top: 40px;
            margin-bottom: 20px;
            font-weight: 600;
        }
        
        h3 {
            font-size: 1.5em;
            color: #424242;
            margin-top: 25px;
            margin-bottom: 15px;
        }

        /* Textboxes and Inputs */
        .gr-textbox textarea, .gr-textbox input {
            border: 1px solid #bdbdbd;
            border-radius: 8px;
            padding: 10px 15px;
            font-size: 1.1em;
            box-shadow: inset 0 1px 3px rgba(0, 0, 0, 0.05);
        }
        .gr-textbox label {
            font-weight: 600;
            color: #555;
            margin-bottom: 8px;
        }

        /* Buttons */
        .gr-button {
            border-radius: 8px;
            padding: 12px 25px;
            font-size: 1.1em;
            font-weight: 600;
            transition: all 0.3s ease;
        }
        .gr-button.primary {
            background-color: #4CAF50; /* Green */
            color: white;
            border: none;
        }
        .gr-button.primary:hover {
            background-color: #43a047;
            box-shadow: 0 2px 8px rgba(0,0,0,0.2);
        }
        .gr-button.secondary {
            background-color: #e3f2fd; /* Light blue background for secondary button */
            color: #424242;
            border: 1px solid #90caf9; /* Light blue border */
        }
        .gr-button.secondary:hover {
            background-color: #bbdefb;
            box-shadow: 0 2px 8px rgba(0,0,0,0.1);
        }
        .gr-button.download {
            background-color: #2196f3; /* Blue for download */
            color: white;
            border: none;
        }
        .gr-button.download:hover {
            background-color: #1976d2;
            box-shadow: 0 2px 8px rgba(0,0,0,0.2);
        }

        /* Markdown Output */
        .gr-markdown {
            background-color: #f9f9f9;
            border: 1px solid #e0e0e0;
            border-radius: 8px;
            padding: 20px;
            line-height: 1.6;
            color: #424242;
            white-space: normal !important; /* Ensure text wraps */
            word-wrap: break-word !important; /* Ensure long words break */
        }
        .gr-markdown p {
            margin-bottom: 10px;
        }
        .gr-markdown ul {
            list-style-type: disc;
            margin-left: 20px;
            padding-left: 0;
        }
        .gr-markdown li {
            margin-bottom: 5px;
        }

        /* Plots */
        .gr-plot {
            border: 1px solid #e0e0e0;
            border-radius: 8px;
            padding: 10px;
            background-color: #ffffff;
            margin-top: 20px;
        }
        
        /* Specific element adjustments */
        #connection-status-textbox {
            font-weight: 500;
            color: #3f51b5;
        }
        #loading-status-textbox {
            font-style: italic;
            color: #757575;
        }
        #input-button-column {
            background-color: #ffffff !important;
        }
        .investor-note {
            color: #2e7d32; /* Green color for investor notes */
            font-style: italic;
            font-weight: 500;
            margin-top: 8px;
            margin-bottom: 8px;
            padding-left: 8px;
            border-left: 2px solid #66bb6a; /* Small green bar on the left */
        }
        """

        with gr.Blocks(theme=gr.themes.Soft(), css=custom_css) as app:
            gr.Markdown("# 📈 BuffetBot - AI Stock Advisor")
            
            # Connection Status
            with gr.Column():
                connection_btn = gr.Button("🔗 Test MCP Connection", variant="secondary")
                connection_status = gr.Textbox(label="MCP Connection Status", interactive=False, elem_id="connection-status-textbox")
                # New vLLM Connection Test
                vllm_connection_btn = gr.Button("🧠 Test vLLM Connection", variant="secondary")
                vllm_connection_status = gr.Textbox(label="vLLM Connection Status", interactive=False, elem_id="vllm-connection-status-textbox")

            with gr.Row():
                with gr.Column(scale=2, elem_id="input-button-column"):
                    ticker_input = gr.Textbox(
                        label="Enter Stock Ticker",
                        placeholder="e.g., AAPL",
                        max_lines=1
                    )
                    generate_btn = gr.Button("Generate Report", variant="primary")
                    loading_status = gr.Textbox(label="Status", interactive=False, visible=False, elem_id="loading-status-textbox")
                
            # Results display components
            output = gr.Markdown(label="Analysis", visible=False)
            revenue_plot = gr.Plot(label="Revenue Growth", visible=False)
            fcf_plot = gr.Plot(label="Free Cash Flow per Share", visible=False)
            shares_plot = gr.Plot(label="Shares Outstanding", visible=False)

            # Static Chart Insights Section
            chart_insights_text = """
            ## Chart Insights
            ### Revenue Growth
            <p class="investor-note">*What to look for: We are looking for companies that consistently grow their revenue year after year, ideally. Consistent growth indicates market acceptance and business expansion.*</p>

            ### Free Cash Flow per Share
            <p class="investor-note">*What to look for: Look for companies with consistently high and growing free cash flow. High FCF indicates a company has strong financial health and flexibility for reinvestment, debt reduction, or shareholder returns.*</p>

            ### Shares Outstanding
            <p class="investor-note">*What to look for: Ideally, look for a declining trend in shares outstanding. This suggests the company is buying back its own shares, which can increase shareholder value by reducing the number of shares in circulation.*</p>
            """
            chart_insights = gr.Markdown(value=chart_insights_text, visible=False, label="Chart Insights")

            download_button = gr.DownloadButton(
                label="Download Analysis as TXT",
                visible=False,
                variant="secondary"
            )
            
            # Hidden state to store analysis text for download
            analysis_text_state = gr.State()

            # Event handlers
            generate_btn.click(
                fn=self.generate_report,
                inputs=[ticker_input],
                outputs=[
                    output,
                    revenue_plot,
                    fcf_plot,
                    shares_plot,
                    analysis_text_state,
                    download_button,
                    loading_status,
                    chart_insights
                ]
            )
            
            ticker_input.submit(
                fn=self.generate_report,
                inputs=[ticker_input],
                outputs=[
                    output,
                    revenue_plot,
                    fcf_plot,
                    shares_plot,
                    analysis_text_state,
                    download_button,
                    loading_status,
                    chart_insights
                ]
            )

            download_button.click(
                fn=self._save_and_return_analysis_file,
                inputs=[analysis_text_state, ticker_input],
                outputs=[download_button],
                show_progress="hidden"
            )

            # Event handler for connection button
            connection_btn.click(
                fn=self._test_mcp_connection,
                inputs=[],
                outputs=[connection_status]
            )
            
            vllm_connection_btn.click(
                fn=self._test_vllm_connection,
                inputs=[],
                outputs=[vllm_connection_status]
            )
        
        return app
    
    def _save_and_return_analysis_file(self, analysis_text: str, ticker: str):
        """Saves the analysis text to a file and returns the path for download."""
        if not analysis_text:
            return None # No file to download if no analysis text

        # Ensure reports directory exists
        reports_dir = "reports"
        os.makedirs(reports_dir, exist_ok=True)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        file_path = os.path.join(reports_dir, f"{ticker}_analysis_{timestamp}.txt")
        
        with open(file_path, "w") as f:
            f.write(analysis_text)
        
        return file_path

    def generate_report(self, ticker):
        """Generate stock analysis report."""
        if not ticker or not ticker.strip():
            return (
                gr.update(value="Error: Please enter a valid stock ticker", visible=True),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                None,
                gr.update(visible=False),
                gr.update(value="❌ Error: Please enter a valid stock ticker", visible=True),
                gr.update(visible=False)
            )

        try:
            # Show loading state
            yield (
                gr.update(value="⏳ Generating report... Please wait.", visible=True),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                None,
                gr.update(visible=False),
                gr.update(value="⏳ Analyzing stock data and generating insights...", visible=True),
                gr.update(visible=False)
            )

            # Call the Modal MCP server for analysis
            headers = {'Content-Type': 'application/json'}
            payload = {"ticker": ticker}
            response = requests.post(f"{self.mcp_server_url}/analyze", headers=headers, json=payload)
            response.raise_for_status()
            analysis_data = response.json()

            # Extract data from the response
            analysis_results = {
                "analysis": analysis_data["analysis"],
                "revenue_chart": pio.from_json(analysis_data["revenue_chart"]) if analysis_data["revenue_chart"] else None,
                "fcf_chart": pio.from_json(analysis_data["fcf_chart"]) if analysis_data["fcf_chart"] else None,
                "shares_chart": pio.from_json(analysis_data["shares_chart"]) if analysis_data["shares_chart"] else None
            }

            # Format the analysis text with a timestamp
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
            formatted_analysis = f"## Analysis Report for {ticker.upper()}\n*Generated on {timestamp}*\n\n{analysis_results['analysis']}"
            
            yield (
                gr.update(value=formatted_analysis, visible=True),
                gr.update(value=analysis_results["revenue_chart"], visible=True),
                gr.update(value=analysis_results["fcf_chart"], visible=True),
                gr.update(value=analysis_results["shares_chart"], visible=True),
                formatted_analysis,
                gr.update(label=f"Download {ticker.upper()} Report", visible=True),
                gr.update(value="✅ Analysis complete!", visible=True),
                gr.update(visible=True)
            )
        except requests.exceptions.RequestException as e:
            error_msg = f"Error analyzing stock: {str(e)}"
            if "429" in str(e):
                error_msg = "Rate limit exceeded. Please try again in a few minutes."
            elif "500" in str(e):
                error_msg = "Server error. Please try again later."
            yield (
                gr.update(value=error_msg, visible=True),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                gr.update(value=None, visible=False),
                None,
                gr.update(visible=False),
                gr.update(value=f"❌ {error_msg}", visible=True),
                gr.update(visible=False)
            )

    def _test_mcp_connection(self):
        """Test connection to Modal MCP server health endpoint."""
        try:
            response = requests.get(f"{self.mcp_server_url}/health", timeout=10)
            if response.status_code == 200:
                return f"✅ Connected to Modal MCP Server: {response.json().get('status', 'OK')}"
            else:
                return f"❌ Connection failed (HTTP {response.status_code}): {response.text}"
        except requests.exceptions.RequestException as e:
            return f"❌ Connection error: {str(e)}"

    def _test_vllm_connection(self):
        """Tests the connection to the Modal vLLM service via the MCP server."""
        try:
            response = requests.get(f"{self.mcp_server_url}/test_llm_connection")
            response.raise_for_status()  # Raise an HTTPError for bad responses (4xx or 5xx)
            status_data = response.json()
            return f"Status: {status_data.get('status', 'Unknown')}"
        except requests.exceptions.RequestException as e:
            return f"Error calling vLLM test endpoint: {e}"

def main():
    """Main entry point for the application."""
    app_instance = BuffetBotApp()
    app = app_instance.create_ui()
    app.launch(server_name="0.0.0.0", server_port=7860, mcp_server=True)

if __name__ == "__main__":
    main()