# app.py
import gradio as gr
from huggingface_hub import InferenceClient
import PyPDF2
import io

# How many characters to include per uploaded PDF (avoid huge inputs / token blowup)
MAX_CHARS_PER_PDF = 20000

def extract_text_from_pdf_fileobj(file_obj) -> str:
    """
    Extract text from a file-like object containing a PDF (Gradio gives a tempfile path
    accessible via file.name, but file might also be an in-memory BytesIO).
    Returns the first MAX_CHARS_PER_PDF characters.
    """
    try:
        # file_obj may be a _TemporaryFileWrapper with .name path or a dict/file-like; handle both
        if hasattr(file_obj, "name"):
            reader = PyPDF2.PdfReader(file_obj.name)
        else:
            # try to read bytes
            file_obj.seek(0)
            reader = PyPDF2.PdfReader(io.BytesIO(file_obj.read()))
    except Exception as e:
        # best-effort: return empty string if cannot parse
        return f"[Could not read PDF: {e}]"

    text_parts = []
    try:
        for page in reader.pages:
            try:
                text_parts.append(page.extract_text() or "")
            except Exception:
                # ignore pages that fail
                continue
    except Exception:
        # some PDFs may throw during iteration — fallback
        pass

    combined = "\n".join(text_parts)
    if len(combined) > MAX_CHARS_PER_PDF:
        combined = combined[:MAX_CHARS_PER_PDF] + "\n\n[TRUNCATED]"
    return combined

def build_context_from_uploaded_files(uploaded_files):
    """
    Given the list returned by gr.File (list of file-like objects), return a single string
    that summarizes / contains the extracted text to be sent as extra context.
    """
    if not uploaded_files:
        return ""

    ctx_parts = []
    for f in uploaded_files:
        try:
            # `f` is a TemporaryFile object from gradio with .name attribute
            extracted = extract_text_from_pdf_fileobj(f)
            header = f"--- Begin extracted text from uploaded file: {getattr(f, 'name', 'uploaded_pdf')} ---\n"
            footer = f"\n--- End of {getattr(f, 'name', 'uploaded_pdf')} ---\n\n"
            ctx_parts.append(header + extracted + footer)
        except Exception as e:
            ctx_parts.append(f"[Error extracting {getattr(f, 'name', 'uploaded_pdf')}: {e}]")

    return "\n".join(ctx_parts)


def respond(
    message,
    history: list[dict[str, str]],
    uploaded_files,                 # NEW: list of uploaded PDF files
    system_message,
    max_tokens,
    temperature,
    top_p,
    hf_token: gr.OAuthToken,
):
    """
    For more information on `huggingface_hub` Inference API support, please check the docs:
    https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
    """
    # Build an extra 'context' system message from uploaded PDFs (if any)
    uploaded_context = build_context_from_uploaded_files(uploaded_files)
    client = InferenceClient(token=hf_token.token, model="openai/gpt-oss-20b")

    # system message first
    messages = [{"role": "system", "content": system_message}]

    # if we have uploaded PDF content, add it as another system message so the model sees it
    if uploaded_context:
        messages.append({"role": "system", "content": "Context from uploaded PDFs:\n\n" + uploaded_context})

    # replay conversation history (assumed to be a list of role/content dicts)
    messages.extend(history)

    # finally the user message
    messages.append({"role": "user", "content": message})

    response = ""

    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        choices = message.choices
        token = ""
        if len(choices) and choices[0].delta.content:
            token = choices[0].delta.content

        response += token
        yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
chatbot = gr.ChatInterface(
    respond,
    type="messages",
    additional_inputs=[
        # NEW: add a File uploader as the first additional input
        gr.File(label="Upload PDFs (optional)", file_count="multiple", file_types=[".pdf"]),
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)

with gr.Blocks() as demo:
    with gr.Sidebar():
        gr.LoginButton()
    chatbot.render()


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