import gradio as gr
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import plotly.graph_objects as go
from plotly.subplots import make_subplots
import warnings
import yfinance as yf
from typing import Dict, List, Tuple
import atexit

warnings.filterwarnings('ignore')

from models import load_timesfm_model, predict_stock_prices
# Mengimpor modul utils secara keseluruhan
import utils 
from config import DEFAULT_PERIOD, DEFAULT_FORECAST_HORIZON, IDX_MARKET_CONFIG

# Load model at startup
model = None

def load_model():
    """Load the Chronos Pipeline"""
    global model
    if model is None:
        model = load_timesfm_model()
    return model

def get_idx_market_status() -> str:
    """Get formatted IDX market status for display"""
    # Mengakses manager melalui fungsi getter
    manager = utils.get_market_manager()
    status = manager.get_status('IDX_STOCKS')
    status_emoji = "🟢" if status.is_open else "🔴"
    status_text = "BUKA" if status.is_open else "TUTUP"
    
    info_lines = [
        f"**Status Saat Ini:** {status.status_text} (WIB)",
        f"**Waktu Buka/Tutup:** {IDX_MARKET_CONFIG['IDX_STOCKS']['open_time']} - {IDX_MARKET_CONFIG['IDX_STOCKS']['close_time']} WIB",
        f"**Waktu Sampai Tutup:** {status.time_until_close}",
        f"**Hari Trading Berikutnya:** {status.next_trading_day}",
        f"**Waktu Sampai Buka:** {status.time_until_open}",
        f"**Terakhir Diperbarui:** {status.last_updated}"
    ]
    
    return f"## {status_emoji} {status.market_name}: {status_text}\n\n" + "\n".join(info_lines)

def qualify_symbol(symbol: str) -> str:
    """Tambahkan ekstensi pasar IDX (.JK) jika belum ada dan pastikan uppercase."""
    # Hapus spasi dan jadikan huruf besar
    symbol = symbol.strip().upper()
    
    # Periksa apakah sudah diakhiri dengan .JK
    if not symbol.endswith(".JK"):
        symbol += ".JK"
        
    return symbol

def analyze_stock(unqualified_symbol, period, forecast_horizon, use_volume):
    """Main analysis function"""
    
    # 1. QUALIFY SYMBOL
    qualified_symbol = qualify_symbol(unqualified_symbol)
    
    try:
        # Load model if not already loaded
        model = load_model()
        
        # Fetch stock data. This now includes technical indicators.
        stock_data = utils.fetch_stock_data(qualified_symbol, period)
        if stock_data is None or len(stock_data) < 30:
            return None, None, None, get_idx_market_status() + f"\n\nError: Data tidak mencukupi untuk {qualified_symbol}. Coba periode yang lebih panjang atau pastikan simbol benar."
        
        # Prepare data for Chronos (unscaled 'Close' prices)
        timesfm_data, scaler = utils.prepare_timesfm_data(stock_data, use_volume)
        
        # Make predictions
        forecast_prices = predict_stock_prices(model, timesfm_data, forecast_horizon)
        
        # Create dates for forecast
        last_date = stock_data.index[-1]
        forecast_dates = pd.date_range(
            start=last_date + timedelta(days=1),
            periods=forecast_horizon,
            freq='D'
        )
        
        # Calculate last price and technical summary
        last_price = stock_data['Close'].iloc[-1]
        
        # Create forecast plot
        fig = utils.create_forecast_plot(stock_data, forecast_dates, forecast_prices, qualified_symbol)
        
        # --- Create Summary Table ---
        # Get last technical indicator values
        last_rsi = stock_data['RSI'].iloc[-1]
        last_macd = stock_data['MACD'].iloc[-1]
        last_macd_signal = stock_data['MACD_Signal'].iloc[-1]
        
        summary_data = {
            'Metrik': ['Harga Saat Ini', 'Puncak Prediksi', 'Dasar Prediksi', 
                      'Perubahan Prediksi (%)', 'Volatilitas Historis', 'RSI (14 Hari)', 'Sinyal MACD'],
            'Nilai': [
                f"Rp {last_price:,.2f}",
                f"Rp {np.max(forecast_prices):,.2f}",
                f"Rp {np.min(forecast_prices):,.2f}",
                f"{((np.mean(forecast_prices) - last_price) / last_price * 100):.2f}%",
                f"{stock_data['Volatility'].iloc[-1] * np.sqrt(252) * 100:.2f}% (Tahunan)",
                f"{last_rsi:.2f}",
                f"{'Beli' if last_macd > last_macd_signal else 'Jual' if last_macd < last_macd_signal else 'Netral'}"
            ]
        }
        
        summary_df = pd.DataFrame(summary_data)
        
        # --- Get Additional Stock Info ---
        stock_info = utils.get_stock_info(qualified_symbol)
        info_text = f"""
        **{stock_info.get('shortName', qualified_symbol)}**
        
        **Statistik Saat Ini:**
        - Kapitalisasi Pasar: {stock_info.get('marketCap', 'N/A')}
        - Volume Trading: {stock_data['Volume'].iloc[-1]:,.0f}
        - 52W Tertinggi: Rp {stock_info.get('fiftyTwoWeekHigh', 'N/A'):,.0f}
        - 52W Terendah: Rp {stock_info.get('fiftyTwoWeekLow', 'N/A'):,.0f}
        
        **Ringkasan Prediksi:**
        - Periode Prediksi: {forecast_horizon} hari
        - Tren Harapan: {'Naik' if np.mean(forecast_prices) > last_price else 'Turun'}
        - Rentang Harga: Rp {np.min(forecast_prices):,.2f} - Rp {np.max(forecast_prices):,.2f}
        """
        
        return fig, summary_df, info_text, get_idx_market_status()
        
    except Exception as e:
        return None, None, None, get_idx_market_status() + f"\n\nError occurred during analysis for {qualified_symbol}: {str(e)}"

def create_interface():
    """Create the Gradio interface"""
    def cleanup_on_exit():
        # MENGGUNAKAN get_market_manager()
        manager = utils.get_market_manager()
        manager.stop()
        print("Market status manager stopped successfully")
    
    atexit.register(cleanup_on_exit)

    with gr.Blocks(
        title="IDX Stock Price Predictor (Chronos-Bolt)",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
        }
        .plot-container {
            height: 500px !important;
        }
        """
    ) as demo:
        
        gr.Markdown(
            """
            # 🇮🇩 IDX Stock Price Predictor
            
            **Didukung oleh Model Chronos-Bolt (Base) Amazon**
            
            Memprediksi harga saham Bursa Efek Indonesia (IDX) menggunakan model *Time Series Foundation Model* Chronos-Bolt.
            
            **Instruksi:**
            1. **Input kode saham IDX (misal: BBCA)**
            2. Tentukan periode data historis (Hari)
            3. Tentukan horison prediksi (Hari)
            4. Klik "Analisis Saham" untuk melihat prediksi
            """
        )
        
        # --- Market Status Display ---
        market_status_output = gr.Markdown(
            value=get_idx_market_status(),
            label="Status Pasar IDX Saat Ini"
        )
        
        gr.Markdown("---")
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Parameter Analisis")
                
                # Mengganti Dropdown menjadi Textbox untuk input manual
                stock_input = gr.Textbox(
                    value="BBCA",
                    label="Input Kode Saham IDX (misal: BBCA)",
                    info="Otomatis ditambahkan ekstensi pasar '.JK'"
                )
                
                period_slider = gr.Slider(
                    minimum=30,
                    maximum=365,
                    value=DEFAULT_PERIOD,
                    step=1,
                    label="Periode Historis (Hari)",
                    info="Jumlah hari data historis yang akan digunakan"
                )
                
                forecast_slider = gr.Slider(
                    minimum=1,
                    maximum=30,
                    value=DEFAULT_FORECAST_HORIZON,
                    step=1,
                    label="Horison Prediksi (Hari)",
                    info="Jumlah hari ke depan yang akan diprediksi"
                )
                
                volume_checkbox = gr.Checkbox(
                    label="Sertakan Volume dalam Analisis (Eksperimental)",
                    value=False,
                    info="Gunakan volume trading sebagai fitur tambahan (mungkin tidak kompatibel dengan Chronos-Bolt)"
                )
                
                analyze_btn = gr.Button(
                    "🔍 Analisis Saham",
                    variant="primary",
                    size="lg"
                )
                
            with gr.Column(scale=2):
                gr.Markdown("### 📈 Hasil Prediksi")
                
                info_output = gr.Markdown(
                    label="Detail Saham",
                    value="Pilih saham dan klik analisis untuk melihat informasi."
                )
                
                with gr.Tab("Grafik Prediksi"):
                    plot_output = gr.Plot(
                        label="Prediksi Harga",
                        show_label=False
                    )
                
                with gr.Tab("Ringkasan Statistik"):
                    summary_output = gr.DataFrame(
                        label="Ringkasan Prediksi",
                        show_label=False
                    )

        
        # Examples section
        gr.Markdown("### 💡 Contoh Cepat")
        examples = gr.Examples(
            examples=[
                ["BBCA", 90, 7, False],
                ["TLKM", 60, 14, True],
                ["UNVR", 120, 10, False],
                ["BMRI", 180, 5, True],
            ],
            inputs=[stock_input, period_slider, forecast_slider, volume_checkbox],
            outputs=[plot_output, summary_output, info_output, market_status_output],
            fn=analyze_stock,
            cache_examples=False
        )
        
        # Footer
        gr.Markdown(
            """
            ---
            **⚠️ Disclaimer:** Alat ini hanya untuk tujuan edukasi. Prediksi pasar saham tidak pasti dan tidak boleh digunakan sebagai nasihat keuangan. Selalu berkonsultasi dengan penasihat keuangan yang berkualifikasi sebelum membuat keputusan investasi.
            
            **Sumber Data:** Yahoo Finance | **Model:** Amazon Chronos-Bolt-Base
            """
        )
        
        # Event handlers
        analyze_btn.click(
            fn=analyze_stock,
            inputs=[stock_input, period_slider, forecast_slider, volume_checkbox],
            outputs=[plot_output, summary_output, info_output, market_status_output],
            show_progress=True
        )
        
        # Initial load to update market status and load model
        demo.load(
            fn=lambda: (load_model(), get_idx_market_status()), 
            outputs=[gr.State(), market_status_output], 
            show_progress=False
        )
    
    return demo

if __name__ == "__main__":
    import atexit
    
    demo = create_interface()
    # Cleanup atexit is registered inside create_interface
    demo.launch(share=True, server_name="0.0.0.0")