import os
import sys
import subprocess
import re
from collections.abc import Iterator
import gradio as gr
from huggingface_hub import hf_hub_download

# Install llama-cpp-python if not present
try:
    from llama_cpp import Llama
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "llama-cpp-python"])
from llama_cpp import Llama

# Install requests if not present
try:
    import requests
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "requests"])
import requests

# Install pypdf if not present
try:
    from pypdf import PdfReader
except ModuleNotFoundError:
    subprocess.check_call([sys.executable, "-m", "pip", "install", "pypdf"])
from pypdf import PdfReader

MAX_MAX_NEW_TOKENS = 512  # Adjusted for document summaries
DEFAULT_MAX_NEW_TOKENS = 256
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "4096"))  # Increased for longer context

DESCRIPTION = """
# DocChat: PDF Summarizer and Q&A (CPU-Only, Ultra-Fast Optimization)
This application delivers an interactive chat interface powered by a highly efficient, small AI model adapted for summarizing and answering questions about any publicly accessible PDF document through specialized prompt engineering. Provide a PDF URL in your message (e.g., 'Load this PDF: https://example.com/document.pdf'), and the bot will download, process it, and be ready for summaries or questions.
It ensures rapid, reasoned responses to user queries. Duplicate this Space for customization or queue-free deployment.
🔎 Model details are available at the [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it) repository; pre-trained and instruction-tuned, further adapted here for document summarization.
<p>Running on CPU 🥶 Inference is heavily optimized for responses in under 10 seconds for simple queries, with output limited to 512 tokens maximum. For longer, more complete responses, increase 'Max New Tokens' in Advanced Settings. Brief delays may occur in free-tier environments due to shared resources, but typical generation speeds reach 20-40 tokens per second.</p>
"""

LICENSE = """
<p/>
---
This application employs the [google/gemma-2b-it](https://huggingface.co/google/gemma-2b-it) model, governed by Google's Gemma Terms of Use. Refer to the [model card](https://huggingface.co/google/gemma-2b-it) and [Gemma documentation](https://ai.google.dev/gemma/terms) for details.
"""

# Download the GGUF model file
model_repo = "mradermacher/gemma-2b-it-GGUF"  # Adjusted to match common repo; update if needed
model_filename = "gemma-2b-it.Q4_K_M.gguf"
model_path = hf_hub_download(repo_id=model_repo, filename=model_filename)

# Load the model with optimizations
llm = Llama(
    model_path=model_path,
    n_ctx=8192,  # Increased context window for document content
    n_batch=512,  # Balanced for CPU efficiency
    n_threads=2,  # Set to 2 for Hugging Face free tier (2 vCPUs)
    n_gpu_layers=0  # Enforce CPU-only execution
)

DEFAULT_SYSTEM_PROMPT = """You are DocChat, a knowledgeable AI assistant specializing in summarizing and analyzing PDF documents. Provide accurate, helpful, reasoned, detailed, and comprehensive summaries or answers based on the provided document. Always base responses on the document content. If no document is loaded, ask the user to provide a publicly accessible PDF URL."""

def process_pdf(url: str) -> str:
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()
        pdf_path = "temp.pdf"
        with open(pdf_path, "wb") as f:
            f.write(response.content)
        reader = PdfReader(pdf_path)
        document_text = ""
        for page in reader.pages:
            text = page.extract_text()
            if text:
                document_text += text + "\n\n"
        # Truncate document text if too long
        max_doc_tokens = 6000  # Leave room for prompt and response
        doc_tokens = llm.tokenize(document_text.encode("utf-8"))
        if len(doc_tokens) > max_doc_tokens:
            document_text = llm.detokenize(doc_tokens[:max_doc_tokens]).decode("utf-8", errors="ignore") + "\n[Document truncated for context length]"
        os.remove(pdf_path)
        return document_text
    except Exception as e:
        return f"Error processing PDF: {str(e)}"

def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int = DEFAULT_MAX_NEW_TOKENS,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    if message.lower().strip() == "hi":
        response = "Hello, I am your document summarizer assistant. Provide a publicly accessible PDF URL to get started."
        yield response
        return

    # Build conversation for Gemma prompt format
    conversation = []
    if system_prompt and not chat_history:
        # Prepend system prompt to first user message for Gemma
        first_message = system_prompt + "\n\n" + message
        conversation.append({"role": "user", "content": first_message})
    else:
        if system_prompt:
            conversation.append({"role": "user", "content": system_prompt})
        for user, assistant in chat_history:
            conversation.append({"role": "user", "content": user})
            conversation.append({"role": "model", "content": assistant})
        conversation.append({"role": "user", "content": message})

    # Construct prompt in Gemma format
    prompt_parts = ["<bos>"]  # Start with BOS token for Gemma
    for dialogue in conversation:
        role = "user" if dialogue["role"] == "user" else "model"
        prompt_parts.append(f"<start_of_turn>{role}\n{dialogue['content']}<end_of_turn>\n")
    prompt_parts.append("<start_of_turn>model\n")
    prompt = "".join(prompt_parts)

    # Tokenize and check length
    input_tokens = llm.tokenize(prompt.encode("utf-8"), add_bos=False)  # Gemma handles BOS internally if needed
    if len(input_tokens) > MAX_INPUT_TOKEN_LENGTH:
        yield "Error: Input too long. Please shorten your query or simplify the request."
        return

    # Generate response
    response = ""
    tokens_generated = 0
    stop_appending = False
    for token in llm.generate(
        input_tokens,
        top_k=top_k,
        top_p=top_p,
        temp=temperature,
        repeat_penalty=repetition_penalty,
    ):
        detokenized = llm.detokenize([token]).decode("utf-8", errors="ignore")
        if not stop_appending:
            response += detokenized
            if "<end_of_turn>" in detokenized or token == llm.token_eos():
                stop_appending = True
            response = response.split("<end_of_turn>")[0].strip()  # Clean up
            yield response
        tokens_generated += 1
        if tokens_generated >= max_new_tokens:
            break

# Gradio interface setup
with gr.Blocks(css="""#chatbot {height: 500px;}""") as demo:
    gr.Markdown(DESCRIPTION)
    chatbot = gr.Chatbot(label="DocChat")
    msg = gr.Textbox(label="Provide PDF URL or ask about the PDF", placeholder="e.g., 'Load PDF: https://example.com/document.pdf' or 'Summarize the key points'")
    with gr.Row():
        submit = gr.Button("Submit")
        clear = gr.Button("Clear")
    with gr.Accordion("Advanced Settings", open=False):
        system_prompt = gr.Textbox(label="System Prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6)
        temperature = gr.Slider(label="Temperature", value=0.6, minimum=0.0, maximum=1.0, step=0.05)
        top_p = gr.Slider(label="Top P", value=0.9, minimum=0.0, maximum=1.0, step=0.05)
        top_k = gr.Slider(label="Top K", value=50, minimum=1, maximum=100, step=1)
        repetition_penalty = gr.Slider(label="Repetition Penalty", value=1.2, minimum=1.0, maximum=2.0, step=0.05)
        max_new_tokens = gr.Slider(label="Max New Tokens", value=DEFAULT_MAX_NEW_TOKENS, minimum=1, maximum=MAX_MAX_NEW_TOKENS, step=1)
    gr.Markdown(LICENSE)

    def user(message, history):
        return "", history + [[message, None]]

    def bot(history, sys_prompt, temp, tp, tk, rp, mnt):
        message = history[-1][0]
        history[-1][1] = ""

        # Detect PDF URL in message
        pdf_url_match = re.search(r'(?i)https?://\S+\.pdf', message)
        new_sys_prompt = sys_prompt
        if pdf_url_match:
            url = pdf_url_match.group(0)
            document_text = process_pdf(url)
            if document_text.startswith("Error"):
                history[-1][1] = document_text
                yield history, new_sys_prompt
                return
            new_sys_prompt = f"""You are DocChat, a knowledgeable AI assistant specializing in summarizing and analyzing PDF documents. Provide accurate, helpful, reasoned, detailed, and comprehensive summaries or answers based on the provided document. Always base responses on the document content.

The document is from {url}. Content excerpt:
{document_text}

Example:
User: Summarize the key points of the document.
Assistant: [Your summary based on the content]."""
            history[-1][1] = "PDF loaded and processed successfully. You can now ask questions or request a summary."
            yield history, new_sys_prompt
            return

        # If no PDF loaded and not loading, remind
        if "Content excerpt:" not in sys_prompt:
            history[-1][1] = "Please provide a publicly accessible PDF URL to load a document first."
            yield history, new_sys_prompt
            return

        # Generate response for queries
        for new_text in generate(message, history[:-1], sys_prompt, mnt, temp, tp, tk, rp):
            history[-1][1] = new_text
            yield history, new_sys_prompt

    submit.click(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens], [chatbot, system_prompt]
    )
    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot, [chatbot, system_prompt, temperature, top_p, top_k, repetition_penalty, max_new_tokens], [chatbot, system_prompt]
    )
    clear.click(lambda: None, None, chatbot, queue=False)

demo.queue(max_size=128).launch()