import gradio as gr
import pandas as pd
import plotly.graph_objects as go
import torch
import os
import io
import base64
from PIL import Image
import matplotlib.pyplot as plt

import os
import sys
import subprocess
import shutil

# --- KRONOS SETUP ---
KRONOS_REPO = "https://github.com/shiyu-coder/Kronos.git"
KRONOS_PATH = "/tmp/kronos_repo"

def setup_kronos():
    if os.path.exists(KRONOS_PATH):
        shutil.rmtree(KRONOS_PATH)
    
    print("Cloning Kronos repository...")
    subprocess.run(["git", "clone", "--depth", "1", KRONOS_REPO, KRONOS_PATH])
    sys.path.insert(0, KRONOS_PATH)
    print("Kronos setup complete.")

# Initialize on boot
setup_kronos()

import kronos_logic

def predict_symbol_ui(symbol, lookback, interval):
    if not symbol:
        return "Please enter a stock symbol.", None
    
    try:
        res = kronos_logic.predict_from_symbol(symbol, lookback, interval)
        status = f"✅ {res['symbol']} ({res['outlook']}): {res['message']} | Current: {res['last_price']} | Predicted: {res['predicted_price']}"
        
        # Convert base64 chart back to image
        img_data = base64.b64decode(res['chart'])
        img = Image.open(io.BytesIO(img_data))
        
        return status, [(img, f"{res['symbol']} Forecast")]
    except Exception as e:
        return f"❌ Error: {str(e)}", None

def predict_batch(files):
    if not files:
        return "No files uploaded.", None
    
    results = []
    gallery_images = []
    
    for file in files:
        try:
            with open(file.name, 'r') as f:
                csv_content = f.read()
                
            res = kronos_logic.predict_from_csv(csv_content)
            
            # Store metadata
            results.append(f"✅ {os.path.basename(file.name)}: {res['message']} (P: {res['predicted_price']})")
            
            # Convert base64 chart back to image for Gradio Gallery
            img_data = base64.b64decode(res['chart'])
            img = Image.open(io.BytesIO(img_data))
            gallery_images.append((img, os.path.basename(file.name)))
            
        except Exception as e:
            results.append(f"❌ {os.path.basename(file.name)}: Error - {str(e)}")
            
    return "\n".join(results), gallery_images

# Build UI
with gr.Blocks(theme=gr.themes.Soft(), title="Kronos AI Predictor") as demo:
    gr.Markdown("# 🚀 Kronos AI Stock Predictor")
    gr.Markdown("Leverage foundation models for high-accuracy financial forecasting.")
    
    with gr.Tabs():
        with gr.Tab("Symbol Search"):
            gr.Markdown("### 🔍 Automated Fetch & Prediction")
            gr.Markdown("Provide a symbol and we'll automatically fetch the latest data from Yahoo Finance.")
            
            with gr.Row():
                with gr.Column(scale=2):
                    symbol_input = gr.Textbox(
                        label="Stock Symbol", 
                        placeholder="e.g. AAPL, TSLA, RELIANCE.NS, BTC-USD",
                        info="For Indian stocks use suffix .NS (NSE) or .BO (BSE). For Crypto use -USD."
                    )
                with gr.Column(scale=1):
                    interval_dropdown = gr.Dropdown(
                        choices=["1d", "1wk", "1mo"], 
                        value="1d", 
                        label="Timeframe",
                        info="Frequency of data points"
                    )
            
            with gr.Row():
                lookback_slider = gr.Slider(
                    minimum=50, 
                    maximum=1000, 
                    value=300, 
                    step=10, 
                    label="Lookback Points",
                    info="Number of historical points to feed the AI"
                )
            
            symbol_btn = gr.Button("Fetch & Generate Forecast", variant="primary")

        with gr.Tab("Bulk CSV Upload"):
            gr.Markdown("### 📂 Batch Process Local Data")
            gr.Markdown("Upload one or more OHLC CSV files for parallel processing.")
            
            file_input = gr.File(label="Upload CSVs", file_count="multiple", file_types=[".csv"])
            run_btn = gr.Button("Run Batch Prediction", variant="secondary")

    with gr.Row():
        status_output = gr.Textbox(label="Operational Status", lines=3, interactive=False)
    
    with gr.Row():
        gallery = gr.Gallery(
            label="Visual Forecasts", 
            show_label=True, 
            elem_id="gallery", 
            columns=[1], 
            object_fit="contain", 
            height="auto"
        )

    # Wire up events
    symbol_btn.click(
        predict_symbol_ui, 
        inputs=[symbol_input, lookback_slider, interval_dropdown], 
        outputs=[status_output, gallery]
    )
    
    run_btn.click(
        predict_batch, 
        inputs=[file_input], 
        outputs=[status_output, gallery]
    )

if __name__ == "__main__":
    demo.launch()