import os
import secrets
from urllib.parse import urlencode
import requests

import streamlit as st
from src.inner.main import run
from src.inner.paths import MARKDOWN_OUTPUT_FILE, PLOT_OUTPUT_FILE

st.title("Query2Plot Demo")
st.write("This is a demo for the Query2Plot application.")



# OAUTH Logic Begin
OAUTH_CLIENT_ID     = os.getenv("OAUTH_CLIENT_ID")
OAUTH_CLIENT_SECRET = os.getenv("OAUTH_CLIENT_SECRET")
SPACE_HOST          = os.getenv("SPACE_HOST")
REDIRECT_URI        = f"https://{SPACE_HOST}"          # must match exactly
AUTH_URL            = "https://huggingface.co/oauth/authorize"
TOKEN_URL           = "https://huggingface.co/oauth/token"
OPENID_PROVIDER_URL = os.getenv("OPENID_PROVIDER_URL", "https://huggingface.co")
SCOPES              = os.getenv("OAUTH_SCOPES", "openid profile")


def oauth_login_url(state: str) -> str:
    query = urlencode({
        "client_id":     OAUTH_CLIENT_ID,
        "redirect_uri":  REDIRECT_URI,
        "response_type": "code",
        "scope":         SCOPES,
        "state":         state,
    })
    return f"{AUTH_URL}?{query}"


def exchange_code(code: str) -> dict:
    r = requests.post(
        TOKEN_URL,
        data={
            "grant_type":   "authorization_code",
            "code":         code,
            "redirect_uri": REDIRECT_URI,
            "client_id":    OAUTH_CLIENT_ID,
        },
        auth=(OAUTH_CLIENT_ID, OAUTH_CLIENT_SECRET),
        timeout=20,
    )
    r.raise_for_status()
    return r.json()


def get_userinfo(access_token: str) -> dict:
    meta = requests.get(
        f"{OPENID_PROVIDER_URL}/.well-known/openid-configuration",
        timeout=20,
    ).json()
    r = requests.get(
        meta["userinfo_endpoint"],
        headers={"Authorization": f"Bearer {access_token}"},
        timeout=20,
    )
    r.raise_for_status()
    return r.json()


# Handle OAUTH Callback
code  = st.query_params.get("code")
state = st.query_params.get("state")

if code and "user" not in st.session_state:
    # NOTE: Because Streamlit resets session_state on redirect, we skip
    # strict state verification here.  If you need it, encode state in the
    # redirect_uri as a path/query param instead.
    try:
        tokens = exchange_code(code)
        user   = get_userinfo(tokens["access_token"])
        st.session_state["tokens"] = tokens
        st.session_state["user"]   = user
    except Exception as e:
        st.error(f"Login failed: {e}")
        st.stop()
    st.query_params.clear()   # remove code/state from URL
    st.rerun()

# Login gate
if "user" not in st.session_state:
    state     = secrets.token_urlsafe(24)
    login_url = oauth_login_url(state)
    # Use target="_blank" so the sign-in page opens in a new tab;
    # this avoids cookie issues when the Space runs inside an iframe.
    st.markdown(
        f'<a href="{login_url}" target="_blank">Sign in with Hugging Face</a>',
        unsafe_allow_html=True,
    )
    st.stop()

st.success(f"Signed in as @{st.session_state['user'].get('preferred_username', 'unknown')}")

# OAUTH Logic End





if "temp_file_path" not in st.session_state: # Ensures temp_file_path is only set once
    st.session_state.temp_file_path = None

uploaded_file = st.file_uploader("Upload a CSV file", type="csv")
if uploaded_file is not None and st.session_state.temp_file_path is None:
    import tempfile
    with tempfile.NamedTemporaryFile(mode='wb', delete=False, suffix=".csv") as temp_file:
        temp_file.write(uploaded_file.getvalue())
        st.session_state.temp_file_path = temp_file.name
    
    st.success(f"File uploaded and saved as {st.session_state.temp_file_path}")

#st.write(f"Enter your query below.")
question = st.text_input("Enter your query:")

if st.button("Run Query2Plot"):
    with st.spinner("Running the program..."):
        st.write('\n')

        question = question + f"\nThe file path to use is: {st.session_state.temp_file_path}"
        result = run(question)

    if MARKDOWN_OUTPUT_FILE.exists():
        with MARKDOWN_OUTPUT_FILE.open("r", encoding="utf-8") as f:
            code_and_output = f.read()
            st.markdown(code_and_output)
    else:
        st.error(f"Markdown output file not found at {MARKDOWN_OUTPUT_FILE}")

    # TO DO: Display the generated plot if it exists
    if PLOT_OUTPUT_FILE.exists():
        st.image(str(PLOT_OUTPUT_FILE), caption="Generated Figure")