import shiny
from shiny import App, ui, reactive
import pandas as pd
import matplotlib.pyplot as plt
import ollama
import io
from shinywidgets import render_plotly
import plotly.express as px
import os

# Fix for Matplotlib permission error
os.environ["MPLCONFIGDIR"] = "/tmp/matplotlib"

# UI
app_ui = ui.page_fluid(
    ui.panel_title("LLM-Powered Data Analytics and Visualization"),
    ui.layout_sidebar(
        ui.sidebar(
            ui.input_file("file", "Upload CSV/Parquet", accept=[".csv", ".parquet"]),
            ui.input_text("question", "Ask about the data:"),
            ui.input_action_button("analyze", "Analyze & Visualize"),
        ),
        ui.layout_columns(
            ui.card(
                ui.output_text("llm_response")
            ),
            ui.card(
                ui.output_plot("plot")
            )
        )
    )
)

# SERVER
def server(input, output, session):
    data = reactive.Value(None)
    
    @reactive.effect
    @reactive.event(input.file)
    def load_data():
        file_info = input.file()
        if file_info is not None:
            ext = file_info["name"].split(".")[-1]
            if ext == "csv":
                df = pd.read_csv(file_info["datapath"])
            elif ext == "parquet":
                df = pd.read_parquet(file_info["datapath"])
            data.set(df)
    
    @output
    @render.text
    @reactive.event(input.analyze)
    def llm_response():
        df = data.get()
        if df is None:
            return "Please upload a dataset."
        question = input.question()
        prompt = f"Dataset: {df.head().to_string()}\nUser Question: {question}\nProvide insights based on the data."
        response = ollama.chat(model="mistral", messages=[{"role": "user", "content": prompt}])
        return response["message"]["content"]
    
    @output
    @render_plotly
    @reactive.event(input.analyze)
    def plot():
        df = data.get()
        if df is None:
            return None
        if len(df.columns) < 2:
            return None
        fig = px.scatter(df, x=df.columns[0], y=df.columns[1])
        return fig

app = App(app_ui, server)