import streamlit as st
import requests
import pandas as pd
import json

st.set_page_config(page_title="MMM Platform", layout="wide")

st.title("MMM Platform")

# Upload CSV file
uploaded_file = st.file_uploader("Upload your CSV file", type=["csv"])

if uploaded_file is not None:
    df = pd.read_csv(uploaded_file)
    st.write("### Preview of Uploaded Data")
    st.dataframe(df.head())

    if "DATE" in df.columns:
        # Convert date column to datetime
        df["DATE"] = pd.to_datetime(df["DATE"], errors="coerce")
        min_date = df["DATE"].min()
        max_date = df["DATE"].max()

        # Date input selectors
        start_date = st.date_input("Select Start Date", min_value=min_date, max_value=max_date, value=min_date)
        end_date = st.date_input("Select End Date", min_value=min_date, max_value=max_date, value=max_date)

        if st.button("Send"):
            try:
                # Prepare data for request
                files = {"file": uploaded_file.getvalue()}
                data = {"start_date": start_date.strftime("%Y-%m-%d"), "end_date": end_date.strftime("%Y-%m-%d")}

                with st.spinner("Processing CSV..."):
                    # Send file and date range to FastAPI endpoint
                    response = requests.post("https://yuvabe-ai-robyn-backend.hf.space/upload_csv/", files=files, data=data)
                
                if response.status_code == 200:
                    response_data = response.json()
                    st.success("CSV processed successfully!")
                    st.json(response_data)
                    
                    st.write("### Robyn Model is Training...")

                    with st.spinner("Training Robyn model..."):
                        robyn_response = requests.post(
                            "https://yuvabe-ai-robyn-backend.hf.space/run_robyn/",
                            files=files,
                            data={"data": json.dumps(response_data["mmm_data_spec"])}
                        )

                    if robyn_response.status_code == 200:
                        robyn_data = robyn_response.json()
                        st.write(robyn_data)
                        report_url = robyn_data.get("report_urls")
                        st.write(report_url)
                        if report_url:
                            st.success("Robyn model training completed successfully!")

                            st.image(report_url)

                        else:
                            st.error("Failed to retrieve report URL.")
                    else:
                        st.error(f"Robyn training failed: {robyn_response.json().get('error', 'Unknown error')}")
                else:
                    st.error(f"CSV processing failed: {response.json().get('error', 'Unknown error')}")

            except Exception as e:
                st.error(f"An error occurred: {e}")

    else:
        st.error("The uploaded CSV must contain a 'DATE' column.")