import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForNextSentencePrediction
import nltk
from typing import List, Tuple
import re
import json
import os

# Download NLTK data
try:
    nltk.data.find('tokenizers/punkt')
except LookupError:
    nltk.download('punkt')

# Page configuration
st.set_page_config(
    page_title="NSP-CitiLink Demo",
    page_icon="📄",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Custom CSS
st.markdown("""
<style>
    .main-header {
        font-size: 2.5rem;
        font-weight: bold;
        color: #1f77b4;
        text-align: center;
        margin-bottom: 1rem;
    }
    .segment-box {
        padding: 1rem;
        margin: 0.5rem 0;
        border-radius: 0.5rem;
        border-left: 4px solid #1f77b4;
        background-color: #f0f2f6;
        color: #1e1e1e;
    }
    /* Dark mode support for segment box */
    @media (prefers-color-scheme: dark) {
        .segment-box {
            background-color: rgb(38, 39, 48);
            color: #e0e0e0;
        }
    }
    .boundary-marker {
        color: #ff4b4b;
        font-weight: bold;
        font-size: 1.2rem;
        text-align: center;
        margin: 1rem 0;
    }
    .metric-box {
        background-color: #e8f4f8;
        padding: 1rem;
        border-radius: 0.5rem;
        text-align: center;
    }
</style>
""", unsafe_allow_html=True)

# Cache model loading
@st.cache_resource
def load_model(model_name):
    """Load the NSP model and tokenizer with fallback."""
    try:
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForNextSentencePrediction.from_pretrained(model_name)
        model.eval()
        return tokenizer, model, None
    except Exception as e:
        error_msg = str(e)
        # If 401 error, likely the model is still being processed by HuggingFace
        if "401" in error_msg or "Unauthorized" in error_msg:
            return None, None, f"⚠️ Model '{model_name}' is still being processed by Hugging Face. Please wait a few minutes and refresh the page."
        else:
            return None, None, f"❌ Error loading model: {error_msg}"

# Load example texts from JSON file
@st.cache_data
def load_example_texts():
    """Load example texts from JSON file."""
    json_path = os.path.join(os.path.dirname(__file__), 'example_texts.json')
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
            # Map the old JSON keys to the new UI names to prevent KeyErrors
            return {
                "Texto Personalizado": "",
                "Atas de Reunião (Português)": data.get("Portuguese Meeting Minutes", "")
            }
    except FileNotFoundError:
        return {
            "Texto Personalizado": "",
            "Atas de Reunião (Português)": ""
        }

def split_into_sentences(text: str) -> List[str]:
    """
    Split text into sentences using NLTK and merge numbered items.
    Uses the same logic as the NSP model for consistency.
    """
    # Use NLTK for sentence splitting
    sentences = nltk.sent_tokenize(text)
    # Clean up and validate sentences
    sentences = [s.strip() for s in sentences if s.strip() and _is_valid_sentence(s.strip())]
    # Merge numbered sentences with following content
    sentences = _merge_numbered_sentences(sentences)
    return sentences

def _is_valid_sentence(sentence: str) -> bool:
    """
    Check if a sentence is valid and should be kept.
    Same validation logic as NSP model.
    """
    if not sentence:
        return False
    
    # Check if it's a numbered item pattern (keep even if short)
    numbered_pattern = r'''
        ^\s*                          # Optional whitespace at start
        (?:                           # Non-capturing group for prefixes
            [-–—]+\s*                 # Dashes (regular, en-dash, em-dash)
            |[•·*+]\s*                # Bullet points
            |§\s*                     # Section symbol
        )?                            # Prefix is optional
        (?:                           # Main numbering patterns
            \d+(?:\.\d+)*[\.\)]\s*    # Numbers with optional decimal parts: "1.", "2.1.", "12.2."
            |[a-z]\)\s*               # Letters with parenthesis: "a)", "b)"
            |[ivxlcdm]+[\.\)]\s*      # Roman numerals: "i.", "ii)", "iv."
            |[-–—]+\s*                # Just dashes: "--", "---"
        )
        $                             # End of string
    '''
    if re.match(numbered_pattern, sentence, re.IGNORECASE | re.VERBOSE):
        return True
    
    # For regular sentences, require at least 3 characters
    return len(sentence) > 3

def _merge_numbered_sentences(sentences: List[str]) -> List[str]:
    """
    Merge sentences that are just numbers/bullets (e.g., "2.", "3.") 
    with the following sentence. Same logic as NSP model.
    """
    if not sentences:
        return sentences
    
    merged_sentences = []
    i = 0
    
    while i < len(sentences):
        current_sentence = sentences[i].strip()
        
        # Check if current sentence is just a numbered item
        numbered_pattern = r'''
            ^\s*                          # Optional whitespace at start
            (?:                           # Non-capturing group for prefixes
                [-–—]+\s*                 # Dashes (regular, en-dash, em-dash)
                |[•·*+]\s*                # Bullet points
                |§\s*                     # Section symbol
            )?                            # Prefix is optional
            (?:                           # Main numbering patterns
                \d+(?:\.\d+)*[\.\)]\s*    # Numbers with optional decimal parts: "1.", "2.1.", "12.2."
                |[a-z]\)\s*               # Letters with parenthesis: "a)", "b)"
                |[ivxlcdm]+[\.\)]\s*      # Roman numerals: "i.", "ii)", "iv."
                |[-–—]+\s*                # Just dashes: "--", "---"
            )
            $                             # End of string
        '''
        
        if re.match(numbered_pattern, current_sentence, re.IGNORECASE | re.VERBOSE):
            # This is a numbered item, find the next non-numbered sentence to merge with
            merge_target_idx = i + 1
            
            # Skip over consecutive numbered items to find actual content
            while (merge_target_idx < len(sentences) and 
                   re.match(numbered_pattern, sentences[merge_target_idx].strip(), re.IGNORECASE | re.VERBOSE)):
                merge_target_idx += 1
            
            if merge_target_idx < len(sentences):
                # Found a non-numbered sentence to merge with
                target_sentence = sentences[merge_target_idx].strip()
                merged_sentence = f"{current_sentence} {target_sentence}"
                merged_sentences.append(merged_sentence)
                i = merge_target_idx + 1  # Skip to after the merged target
            else:
                # No non-numbered sentence found, keep as is
                merged_sentences.append(current_sentence)
                i += 1
        else:
            # Regular sentence, keep as is
            merged_sentences.append(current_sentence)
            i += 1
    
    return merged_sentences

def predict_boundary(sentence_a: str, sentence_b: str, tokenizer, model) -> Tuple[float, float]:
    """
    Predict if there's a topic boundary between two sentences.
    
    Returns:
        Tuple of (is_next_prob, not_next_prob)
    """
    inputs = tokenizer(
        sentence_a, 
        sentence_b, 
        return_tensors="pt",
        truncation=True,
        max_length=512
    )
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=1)
    
    is_next_prob = probs[0][0].item()
    not_next_prob = probs[0][1].item()
    
    return is_next_prob, not_next_prob

def segment_document(text: str, tokenizer, model, threshold: float = 0.5) -> Tuple[List[Tuple[List[str], float]], List[dict]]:
    """
    Segment a document into topics.
    
    Returns:
        Tuple of:
        - List of (segment_sentences, boundary_confidence) tuples
        - List of segment dictionaries with text and offsets
    """
    sentences = split_into_sentences(text)
    
    if len(sentences) == 0:
        return [], []
    
    segments = []
    current_segment = [sentences[0]]
    
    for i in range(1, len(sentences)):
        is_next_prob, not_next_prob = predict_boundary(
            sentences[i-1], 
            sentences[i], 
            tokenizer, 
            model
        )
        
        if not_next_prob > threshold:
            # Topic boundary detected
            segments.append((current_segment, not_next_prob))
            current_segment = [sentences[i]]
        else:
            # Continue current segment
            current_segment.append(sentences[i])
    
    # Add the last segment
    if current_segment:
        segments.append((current_segment, 0.0))  # Last segment has no boundary after it
    
    # Calculate offsets for JSON output
    segments_json = []
    current_offset = 0
    
    for segment_sentences, _ in segments:
        segment_text = ' '.join(segment_sentences)
        
        # Find the actual position in the original text
        start_pos = text.find(segment_sentences[0], current_offset)
        if start_pos == -1:
            # Fallback: use current offset
            start_pos = current_offset
        
        # Find end position after the last sentence in segment
        last_sentence = segment_sentences[-1]
        end_pos = text.find(last_sentence, start_pos)
        if end_pos != -1:
            end_pos += len(last_sentence)
        else:
            # Fallback: estimate based on segment text length
            end_pos = start_pos + len(segment_text)
        
        segments_json.append({
            "text": segment_text,
            "start": start_pos,
            "end": end_pos
        })
        
        current_offset = end_pos
    
    return segments, segments_json

# Main app
def main():
    # Header
    st.markdown('<p class="main-header">📄 NSP-CitiLink: Text Segmentation Demo</p>', unsafe_allow_html=True)
    st.markdown("""
    <p style="text-align: center; color: #666;">
    Automatic text segmentation for city council minutes and administrative documents
    </p>
    """, unsafe_allow_html=True)
    
    # Sidebar
    st.sidebar.header("⚙️ Configuração")

    # Simplificar a seleção de exemplos (Removido os ingleses)
    example = st.sidebar.selectbox(
        "Escolha um exemplo ou insira o seu próprio texto:",
        [
            "Texto Personalizado",
            "Atas de Reunião (Português)"
        ]
    )

    # Novo mapeamento de modelos e localização
    model_configs = {
        "Texto Personalizado": {
            "model": "inesctec/Citilink-NSP-Segmentation-pt",
            "default_threshold": 0.5
        },
        "Atas de Reunião (Português)": {
            "model": "inesctec/Citilink-NSP-Segmentation-pt",
            "default_threshold": 0.5
        }
    }
    
    current_config = model_configs[example]
    
    # Load model with error handling
    with st.spinner(f"🔄 Loading model: {current_config['model']}..."):
        tokenizer, model, error = load_model(current_config['model'])
    
    # Display error if model failed to load
    if error:
        st.sidebar.error(error)
        if "401" in error or "being processed" in error:
            st.sidebar.info("""
            **💡 Tip:** Newly uploaded models can take 5-15 minutes to be fully processed by Hugging Face.
            
            **Workaround:** You can use the Portuguese model (`nsp-citilink`) for English text - it works reasonably well for both languages.
            """)
        # Use Portuguese model as fallback
        st.sidebar.warning("🔄 Attempting to use fallback model: `nsp-citilink`")
        tokenizer, model, fallback_error = load_model("inesctec/CitiLink-NSP-Segmentation-pt")
        if fallback_error:
            st.error("❌ Fallback model also failed to load. Please try again later.")
            st.stop()
    
    threshold = st.sidebar.slider(
        "Boundary Detection Threshold",
        min_value=0.0,
        max_value=1.0,
        value=current_config['default_threshold'],
        step=0.05,
        help="Higher threshold = fewer boundaries (more conservative)"
    )
    
    st.sidebar.markdown("---")
    st.sidebar.markdown("### 📊 About")

    # Show current model info
    current_model_display = "Citilink_NSP_Segmentation"

    st.sidebar.info(f"""
    **Citilink NSP** uses Next Sentence Prediction to identify topic boundaries in administrative documents.

    - **Current Model**: {current_model_display}
    - **Language**: Portuguese
    """)

    st.sidebar.markdown("---")
    st.sidebar.markdown("### 🔗 Resources")
    st.sidebar.markdown("""
    - [📖 Citilink Model Card](https://huggingface.co/inesctec/CitiLink-NSP-Segmentation-pt)
    - [💻 GitHub Repository](https://github.com/jmisidro/segnsp)
    - [📊 CitiLink Dataset](https://github.com/INESCTEC/citilink-dataset)
    """)
    
    # Main content
    col1, col2 = st.columns([1, 1])
    
    with col1:
        st.subheader("📝 Input Document")
        
        # Load example texts from JSON
        example_texts = load_example_texts()
        
        if example == "Custom Text":
            input_text = st.text_area(
                "Enter your text here:",
                height=400,
                placeholder="Paste your document text here..."
            )
        else:
            input_text = st.text_area(
                f"Example: {example}",
                value=example_texts[example],
                height=400
            )
        
        segment_button = st.button("🔍 Segment Document", type="primary", use_container_width=True) 
            
    with col2:
        st.subheader("📊 Segmentation Results")
        
        if segment_button and input_text:
            with st.spinner("🔄 Analyzing document..."):
                # Perform segmentation
                segments, segments_json = segment_document(input_text, tokenizer, model, threshold)
                
                if not segments:
                    st.warning("⚠️ No text to segment. Please enter some text.")
                else:
                    # Display metrics
                    metric_col1, metric_col2, metric_col3 = st.columns(3)
                    
                    with metric_col1:
                        st.markdown('<div class="metric-box">', unsafe_allow_html=True)
                        st.metric("Total Segments", len(segments))
                        st.markdown('</div>', unsafe_allow_html=True)
                    
                    with metric_col2:
                        total_sentences = sum(len(seg[0]) for seg in segments)
                        st.markdown('<div class="metric-box">', unsafe_allow_html=True)
                        st.metric("Total Sentences", total_sentences)
                        st.markdown('</div>', unsafe_allow_html=True)
                    
                    with metric_col3:
                        avg_segment_size = total_sentences / len(segments) if segments else 0
                        st.markdown('<div class="metric-box">', unsafe_allow_html=True)
                        st.metric("Avg. Sentences/Segment", f"{avg_segment_size:.1f}")
                        st.markdown('</div>', unsafe_allow_html=True)
                    
                    st.markdown("---")
                    
                    # Display segments
                    st.markdown("### 📑 Segmented Document")
                    
                    for idx, (segment_sentences, boundary_conf) in enumerate(segments, 1):
                        # Display segment
                        st.markdown(f"""
                        <div class="segment-box">
                            <strong>📌 Segment {idx}</strong> ({len(segment_sentences)} sentences)
                            <br><br>
                            {' '.join(segment_sentences)}
                        </div>
                        """, unsafe_allow_html=True)
                        
                        # Display boundary marker if not last segment
                        if idx < len(segments):
                            confidence_pct = boundary_conf * 100
                            st.markdown(f"""
                            <div class="boundary-marker">
                                🔴 Topic Boundary Detected (confidence: {confidence_pct:.1f}%)
                            </div>
                            """, unsafe_allow_html=True)
                    
                    # Download buttons and JSON view
                    st.markdown("---")
                    
                    # Create text output
                    result_text = ""
                    for idx, (segment_sentences, _) in enumerate(segments, 1):
                        result_text += f"=== SEGMENT {idx} ===\n"
                        result_text += ' '.join(segment_sentences)
                        result_text += "\n\n"
                    
                    # Create JSON output
                    json_output = {
                        "segments": segments_json
                    }
                    json_str = json.dumps(json_output, indent=2, ensure_ascii=False)
                    
                    # Download and view options
                    col_download1, col_download2 = st.columns(2)
                    
                    with col_download1:
                        st.download_button(
                            label="⬇️ Download Segmented Text",
                            data=result_text,
                            file_name="segmented_document.txt",
                            mime="text/plain",
                            use_container_width=True
                        )
                    
                    with col_download2:
                        st.download_button(
                            label="⬇️ Download JSON",
                            data=json_str,
                            file_name="segments.json",
                            mime="application/json",
                            use_container_width=True
                        )
                    
                    # Expandable JSON viewer
                    with st.expander("📋 View Raw JSON (with offsets)"):
                        st.json(json_output)
        else:
            st.info("👈 Enter text in the input box and click 'Segment Document' to begin.")
            
            # Show example predictions
            st.markdown("### 🎯 How It Works")
            st.markdown("""
            The model analyzes pairs of consecutive sentences and predicts:
            - **Same Topic** (label 0): Sentences discuss the same subject
            - **Topic Boundary** (label 1): New topic begins at second sentence
            
            **Example:**
            """)
            
            st.code("""
Sentence A: "By the President, minutes no. 28 of 20.12.2023 were present at the meeting."
Sentence B: "After considering and analyzing the matter, the Municipal Executive unanimously decided to approve minute no. 28 of 12.20.2023."
→ Prediction: Same Topic (confidence: 76%)
Sentence A: "After considering and analyzing the matter, the Municipal Executive unanimously decided to approve minute no. 28 of 12.20.2023."
Sentence B: "There were no various processes and requests to submit."
→ Prediction: Topic Boundary (confidence: 82%)
            """)

if __name__ == "__main__":
    main()