import streamlit as st
import os
import tempfile
from uuid import uuid4
from concurrent.futures import ThreadPoolExecutor, Future

from ingestion import ingest_pdf
from dotenv import load_dotenv, find_dotenv

load_dotenv(find_dotenv(), override=False)
st.set_page_config(page_title="PDF Ingestion", layout="wide")

if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid4())
if "ingestion_future" not in st.session_state:
    st.session_state.ingestion_future = None
if "last_output" not in st.session_state:
    st.session_state.last_output = None

st.title("Ingestion Layer")
st.caption("Upload a PDF to extract, normalize, chunk, and persist for later embedding.")

with st.sidebar:
    st.header("Settings")
    parser = st.selectbox("Parser", ["pypdf", "pdfplumber"], index=0)
    chunk_size = st.number_input("Chunk size (chars)", min_value=500, max_value=10000, value=4000, step=100)
    overlap = st.number_input("Overlap (chars)", min_value=0, max_value=2000, value=400, step=50)

uploaded = st.file_uploader("Upload PDF", type=["pdf"])

def _run_ingestion(tmp_path: str, sid: str, parser: str, chunk_size: int, overlap: int):
    return ingest_pdf(tmp_path, sid, parser=parser, chunk_size=chunk_size, overlap=overlap)

col1, col2 = st.columns([2, 1])
with col1:
    if uploaded is not None and st.button("Start Ingestion", type="primary"):
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded.read())
            tmp_path = tmp.name

        executor = ThreadPoolExecutor(max_workers=1)
        future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, parser, int(chunk_size), int(overlap))
        st.session_state.ingestion_future = future
        st.info("Ingestion started. Processing in background...")

    future: Future = st.session_state.ingestion_future
    if future is not None:
        if future.done():
            try:
                chunks, out_path = future.result()
                st.session_state.last_output = (len(chunks), out_path)
                st.success(f"Ingestion complete: {len(chunks)} chunks written to {out_path}")
                if chunks:
                    st.subheader("Preview (first chunk)")
                    st.code(chunks[0].content[:2000])
            except Exception as e:
                st.error(f"Ingestion failed: {e}")
            finally:
                st.session_state.ingestion_future = None
        else:
            st.status("Processing...", state="running")

with col2:
    st.subheader("Session")
    st.code(st.session_state.session_id)
    if st.session_state.last_output:
        cnt, path = st.session_state.last_output
        st.metric("Chunks", cnt)
        st.write("Output:")
        st.code(path)
    if uploaded is not None and st.button("Retry with fallback parser"):
        fb = "pdfplumber" if parser == "pypdf" else "pypdf"
        with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp:
            tmp.write(uploaded.read())
            tmp_path = tmp.name
        executor = ThreadPoolExecutor(max_workers=1)
        future: Future = executor.submit(_run_ingestion, tmp_path, st.session_state.session_id, fb, int(chunk_size), int(overlap))
        st.session_state.ingestion_future = future
        st.info(f"Retrying with {fb}...")