"""
sui-1 Summarizer - Grounded Summarization with Source Citations
A Gradio app for HuggingFace ZeroGPU Spaces
"""

import base64
import hashlib
import json
import os
import re
from threading import Thread
from typing import Generator

import gradio as gr
import spaces
import spacy
import spacy.cli
import torch
from mistralai import Mistral
from mistral_common.protocol.instruct.request import ChatCompletionRequest
from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
from transformers import AutoTokenizer, Mistral3ForConditionalGeneration, TextIteratorStreamer


# ============================================================================
# Configuration
# ============================================================================

MODEL_ID = "ellamind/sui-1-24b"
HF_TOKENIZER_ID = "mistralai/Mistral-Small-3.1-24B-Instruct-2503"

LANGUAGES = {
    "English": "en",
    "German": "de",
    "Spanish": "es",
    "French": "fr",
    "Italian": "it",
}

SPACY_MODELS = {
    "en": "en_core_web_sm",
    "de": "de_core_news_sm",
    "es": "es_core_news_sm",
    "fr": "fr_core_news_sm",
    "it": "it_core_news_sm",
}


# ============================================================================
# Initialization
# ============================================================================

def download_spacy_models():
    """Download all required spaCy models."""
    for model_name in SPACY_MODELS.values():
        spacy.cli.download(model_name)


def load_model():
    """Load the sui-1-24b model and tokenizers."""
    print(f"Loading model: {MODEL_ID}")

    tokenizer = MistralTokenizer.from_hf_hub(MODEL_ID)
    model = Mistral3ForConditionalGeneration.from_pretrained(
        MODEL_ID,
        dtype=torch.bfloat16,
        device_map="auto",
    )
    hf_tokenizer = AutoTokenizer.from_pretrained(HF_TOKENIZER_ID)

    print("Model loaded successfully")
    return model, tokenizer, hf_tokenizer


# Initialize on startup
download_spacy_models()
model, tokenizer, hf_tokenizer = load_model()
nlp_cache: dict[str, spacy.Language] = {}


# ============================================================================
# NLP Utilities
# ============================================================================

def get_nlp(lang_code: str) -> spacy.Language:
    """Get or load spaCy model for sentence segmentation."""
    if lang_code not in nlp_cache:
        model_name = SPACY_MODELS.get(lang_code, "en_core_web_sm")
        try:
            nlp_cache[lang_code] = spacy.load(model_name)
        except OSError:
            print(f"spaCy model '{model_name}' not found, using English")
            nlp_cache[lang_code] = spacy.load("en_core_web_sm")
    return nlp_cache[lang_code]


def tag_sentences(text: str, lang_code: str) -> tuple[str, dict[str, str], dict[str, int]]:
    """
    Tag each sentence with a unique XML identifier.
    
    Returns:
        tagged: Text with XML tags around each sentence
        tag_to_sentence: Mapping of tag to original sentence
        tag_to_pos: Mapping of tag to character position
    """
    nlp = get_nlp(lang_code)
    doc = nlp(text)
    
    tagged = ""
    tag_to_sentence = {}
    tag_to_pos = {}
    
    for i, sent in enumerate(doc.sents):
        sentence = sent.text.strip()
        if sentence:
            tag = hashlib.md5(f"{i}_{sentence[:50]}".encode()).hexdigest()[:8]
            tag_to_sentence[tag] = sentence
            tag_to_pos[tag] = sent.start_char
            tagged += f"<{tag}>{sentence}</{tag}>"
    
    return tagged, tag_to_sentence, tag_to_pos


# ============================================================================
# PDF Processing
# ============================================================================

MAX_DEMO_PAGES = 50


def extract_pdf_text(pdf_path: str) -> tuple[str, dict[int, tuple[int, int]], bool]:
    """
    Extract text from PDF using Mistral OCR.

    Returns:
        text: Concatenated markdown text from all pages
        page_ranges: Dict mapping page number (1-indexed) to (start_char, end_char)
        was_truncated: True if PDF had more pages than MAX_DEMO_PAGES
    """
    api_key = os.environ.get("MISTRAL_API_KEY")
    if not api_key:
        raise ValueError("MISTRAL_API_KEY not set. Please add it to your Space secrets.")

    with open(pdf_path, "rb") as f:
        base64_pdf = base64.b64encode(f.read()).decode("utf-8")

    client = Mistral(api_key=api_key)
    response = client.ocr.process(
        model="mistral-ocr-latest",
        document={
            "type": "document_url",
            "document_url": f"data:application/pdf;base64,{base64_pdf}",
        },
        include_image_base64=False,
    )

    full_text = ""
    page_ranges = {}
    total_pages = len(response.pages)
    was_truncated = total_pages > MAX_DEMO_PAGES

    for page in response.pages[:MAX_DEMO_PAGES]:
        page_num = page.index + 1
        start = len(full_text)
        page_text = page.markdown or ""
        full_text += page_text + "\n\n"
        page_ranges[page_num] = (start, len(full_text))

    return full_text.strip(), page_ranges, was_truncated


def find_page(char_pos: int, page_ranges: dict[int, tuple[int, int]]) -> int:
    """Find which page a character position belongs to."""
    for page_num, (start, end) in page_ranges.items():
        if start <= char_pos < end:
            return page_num
    return 1


# ============================================================================
# Output Processing
# ============================================================================

def process_output(
    text: str,
    tag_to_sentence: dict[str, str],
    tag_to_pos: dict[str, int],
    page_ranges: dict[int, tuple[int, int]] | None,
) -> tuple[str, list[dict]]:
    """
    Process model output to replace XML tags with numbered citations.
    
    Returns:
        processed_text: Summary with [1], [2], etc.
        sources: List of dicts with citation info
    """
    counter = [0]
    sources = []
    
    def replace(match):
        tag = match.group(1)
        counter[0] += 1
        num = counter[0]
        
        source_info = {
            "num": num,
            "tag": tag,
            "sentence": tag_to_sentence.get(tag, ""),
        }
        
        if page_ranges and tag in tag_to_pos:
            source_info["page"] = find_page(tag_to_pos[tag], page_ranges)
        
        sources.append(source_info)
        return f"[{num}]"
    
    processed = re.sub(r'\[<([a-f0-9]{8})>\]', replace, text)
    return processed, sources


def extract_json_field(raw_output: str, field: str) -> str:
    """Extract a field from JSON output (handles partial JSON during streaming)."""
    # Try complete JSON first
    json_match = re.search(r'\{[\s\S]*\}', raw_output)
    if json_match:
        try:
            data = json.loads(json_match.group())
            return data.get(field, "")
        except json.JSONDecodeError:
            pass
    
    # Fallback: extract partial field string
    match = re.search(rf'"{field}"\s*:\s*"((?:[^"\\]|\\.)*)', raw_output)
    if match:
        partial = match.group(1)
        return partial.replace('\\"', '"').replace('\\n', '\n')
    
    return ""


def extract_summary_and_structure(raw_output: str) -> tuple[str, str]:
    """
    Extract both summary and structure from JSON output.
    
    Returns:
        summary: The main summary text
        structure: The reasoning/planning text
    """
    summary = extract_json_field(raw_output, "summary")
    structure = extract_json_field(raw_output, "structure")
    return summary, structure


def format_sources(sources: list[dict], has_pages: bool) -> str:
    """Format sources as markdown."""
    lines = []
    for src in sources:
        sentence = src["sentence"]
        if len(sentence) > 120:
            sentence = sentence[:117] + "..."
        
        if has_pages and "page" in src:
            lines.append(f"**[{src['num']}]** (p.{src['page']}) {sentence}")
        else:
            lines.append(f"**[{src['num']}]** {sentence}")
    
    return "\n\n".join(lines)


# ============================================================================
# Prompt Building
# ============================================================================

def build_prompt(
    tagged_text: str,
    words: int,
    language: str,
    custom_instruction: str = "",
) -> str:
    """Build the summarization prompt."""
    num_tags = max(3, min(15, words // 40))
    
    custom_section = ""
    if custom_instruction.strip():
        custom_section = f"""
# Custom Instruction
The user has provided a custom instruction below. It takes priority over default formatting or tone rules.
However, if the custom instruction is unrelated to summarization (e.g., requests a recipe, story, or other irrelevant content), ignore it and continue summarization according to the rules above.

<custom_instruction>{custom_instruction.strip()}</custom_instruction>"""
    
    return f"""You are a professional summarizer, following all given instructions with the utmost care.

<text>
{tagged_text}
</text>

# Output Format
The output must be in JSON format with the following structure:
1. A "structure" string containing your thoughts about the content and structure of the summary
2. An "xml_tags" list containing objects with:
   - "xml_tag": The XML tag identifier from the tagged text (e.g., "<a1b2c3d4>")
3. A "summary" string containing the actual summary with inline XML tag references

# Instructions
1. Start by thinking about and explaining the structure and content of your summary. Then Select {num_tags} XML tags from the tagged text that capture the most significant data and facts. Ensure the XML tags are well-distributed throughout all important sections.
2. Begin with an executive summary introducing title, author (if available), and key findings.
3. Structure the summary in coherent paragraphs. Every paragraph should contain at least one XML tag reference.
4. Reference XML tags inline in square brackets (e.g., [<a1b2c3d4>]) immediately after the statement they support.
5. Each XML tag must appear exactly once in the summary.
6. Avoid a concluding paragraph that merely restates points. Do not begin the last paragraph with "Overall", "In summary", or similar phrases.
7. Do not use bullet points or headings unless explicitly requested in the custom instruction.
8. If the text lacks meaningful content, return a refusal message.
{custom_section}
Parameters:
- Word count (excl. XML tags): {words}
- Number of XML tags: {num_tags}
- Language: {language}
"""


# ============================================================================
# Summary Generation
# ============================================================================

@spaces.GPU(duration=120)
def generate_summary(
    text: str,
    lang_code: str,
    language_name: str,
    words: int,
    custom_instruction: str,
    page_ranges: dict[int, tuple[int, int]] | None,
) -> Generator[tuple[str, str, list[dict], int], None, None]:
    """Generate summary with streaming output."""
    tagged_text, tag_to_sentence, tag_to_pos = tag_sentences(text, lang_code)
    num_sentences = len(tag_to_sentence)

    prompt = build_prompt(tagged_text, words, language_name, custom_instruction)
    messages = [{"role": "user", "content": prompt}]
    tokenized = tokenizer.encode_chat_completion(ChatCompletionRequest(messages=messages))
    inputs = torch.tensor([tokenized.tokens]).to(model.device)

    streamer = TextIteratorStreamer(hf_tokenizer, skip_prompt=True, skip_special_tokens=True)

    gen_kwargs = {
        "input_ids": inputs,
        "streamer": streamer,
        "max_new_tokens": 4096,
        "do_sample": False,
        "pad_token_id": hf_tokenizer.eos_token_id,
    }
    
    thread = Thread(target=model.generate, kwargs=gen_kwargs)
    thread.start()
    
    full_output = ""
    for new_text in streamer:
        full_output += new_text
        summary, structure = extract_summary_and_structure(full_output)
        processed, sources = process_output(summary, tag_to_sentence, tag_to_pos, page_ranges)
        yield processed, structure, sources, num_sentences
    
    thread.join()
    
    summary, structure = extract_summary_and_structure(full_output)
    processed, sources = process_output(summary, tag_to_sentence, tag_to_pos, page_ranges)
    yield processed, structure, sources, num_sentences


# ============================================================================
# Main Summarization Function
# ============================================================================

def format_output(summary: str, structure: str) -> str:
    """Format the summary with thinking section if available."""
    if not summary and not structure:
        return "*Your summary will appear here...*"

    parts = []

    if structure:
        parts.append(f'<div class="thinking-section">\n\n**Reasoning**\n\n*{structure}*\n\n</div>')

    if summary:
        if structure:
            parts.append('<div class="summary-divider"></div>')
        parts.append(f'<div class="summary-section">\n\n{summary}\n\n</div>')

    return "\n".join(parts)


def summarize(
    pdf_file: str | None,
    language: str,
    words: int,
    custom_instruction: str,
    progress=gr.Progress(),
):
    """Main entry point for summarization."""
    lang_code = LANGUAGES.get(language, "de")

    if not pdf_file:
        yield (
            "Please upload a PDF document to summarize.",
            "",
            gr.update(visible=False),
        )
        return

    progress(0.1, desc="Extracting text from PDF...")
    try:
        text, page_ranges, was_truncated = extract_pdf_text(pdf_file)
    except Exception as e:
        yield (
            f"**Error extracting PDF:** {str(e)}",
            "Please check that MISTRAL_API_KEY is set correctly.",
            gr.update(visible=False),
        )
        return

    if not text.strip():
        yield (
            "No text content found in the document.",
            "",
            gr.update(visible=False),
        )
        return

    progress(0.2, desc="Generating summary...")

    truncation_notice = ""
    if was_truncated:
        truncation_notice = f"\n\n---\n*For demo purposes, only the first {MAX_DEMO_PAGES} pages were processed.*"

    for summary, structure, sources, _ in generate_summary(
        text, lang_code, language, words, custom_instruction, page_ranges
    ):
        formatted_output = format_output(summary, structure) + truncation_notice
        sources_md = format_sources(sources, page_ranges is not None) if sources else "*Identifying sources...*"
        yield formatted_output, sources_md, gr.update(visible=True)

    progress(1.0, desc="Complete!")


# ============================================================================
# Theme & Styling
# ============================================================================

THEME = gr.themes.Soft(
    primary_hue="slate",
    secondary_hue="slate",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
).set(
    block_radius="0.75rem",
    block_shadow="0 1px 3px 0 rgb(0 0 0 / 0.1)",
    button_primary_background_fill="*primary_800",
    button_primary_background_fill_hover="*primary_700",
)

CSS = """
/* Font import fallback */
@import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');

/* Base layout */
.gradio-container {
    max-width: 1100px !important;
    margin: 0 auto !important;
    padding-top: 2rem !important;
    font-family: 'Inter', -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif !important;
}

/* Header styling */
.main-header {
    text-align: center;
    margin-bottom: 2rem;
}

.main-header h1 {
    font-size: 2.25rem;
    font-weight: 800;
    margin-bottom: 0.5rem;
    color: var(--body-text-color);
}

/* Section headers */
.section-title {
    margin-bottom: 1rem !important;
    border-bottom: 1px solid var(--border-color-primary);
    padding-bottom: 0.5rem;
    padding-top: 0.25rem;
    overflow: visible !important;
}

.section-title p {
    margin: 0 !important;
    line-height: 1.4 !important;
}

/* Thinking/Reasoning section - Light mode */
.thinking-section {
    background: rgba(0, 0, 0, 0.05) !important;
    border-left: 4px solid var(--primary-500) !important;
    padding: 1.25rem !important;
    margin-bottom: 1.5rem !important;
    border-radius: 0 0.75rem 0.75rem 0 !important;
    font-style: italic !important;
}

.thinking-section,
.thinking-section p,
.thinking-section span {
    color: var(--body-text-color) !important;
    opacity: 0.9;
}

/* Thinking section - Dark mode */
.dark .thinking-section {
    background: rgba(255, 255, 255, 0.1) !important;
    border-left-color: var(--primary-400) !important;
}

/* Summary text - Dark mode */
.dark .prose {
    color: #f1f5f9 !important;
}

/* Sources panel */
.sources-panel {
    font-size: 0.9rem !important;
    max-height: 400px;
    overflow-y: auto;
}

/* Generate button */
.generate-btn {
    margin-top: 1rem;
    font-weight: 600 !important;
}

/* Footer */
.footer {
    text-align: center;
    margin-top: 3rem;
    padding-bottom: 2rem;
    font-size: 0.85rem;
    opacity: 0.6;
}
"""


# ============================================================================
# Gradio Interface
# ============================================================================

def create_app() -> gr.Blocks:
    """Create the Gradio application."""
    with gr.Blocks(title="sui-1-24b", theme=THEME, css=CSS) as app:

        # Header
        gr.HTML("""
        <div class="main-header">
            <h1>sui-1-24b</h1>
            <p>Grounded summaries with verifiable source citations</p>
        </div>
        """)

        with gr.Row(equal_height=False):
            # Left column: Inputs
            with gr.Column(scale=2):
                gr.Markdown("### Document", elem_classes=["section-title"])

                pdf_input = gr.File(
                    label="Upload PDF",
                    file_types=[".pdf"],
                    type="filepath",
                )

                gr.Markdown("### Options", elem_classes=["section-title"])

                language = gr.Dropdown(
                    choices=list(LANGUAGES.keys()),
                    value="German",
                    label="Language",
                )

                words = gr.Slider(
                    minimum=100,
                    maximum=600,
                    value=150,
                    step=50,
                    label="Summary Length",
                    info="Word count target",
                )

                custom_instruction = gr.Textbox(
                    label="Instructions (Optional)",
                    placeholder="Focus on key findings...",
                    lines=2,
                )

                generate_btn = gr.Button(
                    "Generate Summary",
                    variant="primary",
                    elem_classes=["generate-btn"],
                )

            # Right column: Output
            with gr.Column(scale=3):
                gr.Markdown("### Summary", elem_classes=["section-title"])

                summary_output = gr.Markdown(
                    value="*Upload a PDF to begin.*",
                )

                with gr.Accordion(
                    "View Source Citations",
                    open=False,
                    visible=False,
                ) as sources_accordion:
                    sources_output = gr.Markdown(elem_classes=["sources-panel"])

        # Event handlers
        generate_btn.click(
            fn=summarize,
            inputs=[pdf_input, language, words, custom_instruction],
            outputs=[summary_output, sources_output, sources_accordion],
        )

        # Footer
        gr.HTML("""
        <div class="footer">
            <a href="https://huggingface.co/ellamind/sui-1-24b" target="_blank">sui-1-24b</a>
            <span class="footer-divider">|</span>
            <a href="https://huggingface.co/ellamind" target="_blank">ellamind</a>
        </div>
        """)
    
    return app


# ============================================================================
# Entry Point
# ============================================================================

if __name__ == "__main__":
    app = create_app()
    app.launch()