import os
import gradio as gr
import pandas as pd
import numpy as np
import re
from huggingface_hub import InferenceClient

def load_data(file_obj):
    """Safely loads CSV, Excel, or TXT file into a Pandas DataFrame."""
    if file_obj is None:
        return None, gr.update(choices=[], visible=False), "Please upload a file."
    
    file_path = file_obj.name
    ext = os.path.splitext(file_path)[1].lower()
    
    try:
        if ext == '.csv':
            df = pd.read_csv(file_path)
        elif ext in ['.xls', '.xlsx']:
            df = pd.read_excel(file_path)
        elif ext == '.txt':
            with open(file_path, 'r', encoding='utf-8') as f:
                content = f.read()
            df = pd.DataFrame({'text': [content]})
        else:
            return None, gr.update(choices=[], visible=False), "Unsupported file format. Please upload .csv, .xlsx, or .txt."
        
        string_cols = [col for col in df.columns if df[col].dtype == 'object' or df[col].astype(str).str.len().mean() > 5]
        if not string_cols:
            string_cols = list(df.columns)
            
        return df, gr.update(choices=string_cols, value=string_cols[0], visible=True), f"Successfully loaded dataset with {len(df)} rows."
    except Exception as e:
        return None, gr.update(choices=[], visible=False), f"Error loading file: {str(e)}"

def run_extraw_cpu(text, ratio=0.3):
    """Rule-based TF-IDF sentence extractive summarizer running entirely locally on CPU."""
    # Split text into sentences
    sentences = re.split(r'(?<=[.!?])\s+', text)
    sentences = [s.strip() for s in sentences if len(s.strip()) > 10]
    
    if len(sentences) <= 3:
        return text  # Too short to summarize
        
    # Calculate word frequencies
    words = re.findall(r'\b[a-zA-Z]{3,}\b', text.lower())
    stopwords = {'the', 'and', 'for', 'that', 'with', 'this', 'have', 'from', 'your', 'will', 'not', 'are', 'was', 'were', 'but', 'how'}
    word_freqs = {}
    for word in words:
        if word not in stopwords:
            word_freqs[word] = word_freqs.get(word, 0) + 1
            
    if not word_freqs:
        return " ".join(sentences[:2])
        
    # Normalize frequencies
    max_freq = max(word_freqs.values())
    for word in word_freqs:
        word_freqs[word] = word_freqs[word] / max_freq
        
    # Score sentences
    sentence_scores = {}
    for i, sent in enumerate(sentences):
        score = 0
        sent_words = re.findall(r'\b[a-zA-Z]{3,}\b', sent.lower())
        for word in sent_words:
            if word in word_freqs:
                score += word_freqs[word]
        sentence_scores[i] = score / max(1, len(sent_words)) # normalize by length to prevent biased long sentences
        
    # Determine number of sentences to extract
    n_sentences = max(1, int(len(sentences) * ratio))
    
    # Select top sentences
    top_indices = sorted(sentence_scores, key=sentence_scores.get, reverse=True)[:n_sentences]
    
    # Sort indices to preserve original chronological order
    top_indices.sort()
    
    summary = " ".join([sentences[idx] for idx in top_indices])
    return summary

def run_transformer_summarize(text, hf_token, model_name, ratio):
    """Summarizes using Hugging Face Serverless Inference API."""
    if not hf_token:
        raise ValueError("Hugging Face API Token is required for Transformer Mode.")
        
    client = InferenceClient(token=hf_token)
    
    # Calculate desired length bounds based on text length
    words_count = len(text.split())
    max_len = max(30, int(words_count * ratio * 1.2))
    min_len = max(10, int(words_count * ratio * 0.8))
    
    try:
        resp = client.summarization(
            text=text,
            model=model_name,
            parameters={"max_length": max_len, "min_length": min_len}
        )
        return resp.get("summary_text", "")
    except Exception as e:
        raise RuntimeError(f"Hugging Face API error: {str(e)}")

def process_summarization(text_input, file_obj, text_col, method, hf_token, hf_model, length_ratio):
    # Parse documents
    docs = []
    if file_obj is not None:
        df, _, _ = load_data(file_obj)
        if df is not None and text_col in df.columns:
            docs = df[text_col].astype(str).fillna("").tolist()
    elif text_input and text_input.strip():
        docs = [text_input]
        
    if not docs:
        return None, None, "Please enter text or upload a valid dataset first."
        
    # Standardize ratio
    ratio_dict = {"Short Summary (15%)": 0.15, "Medium Summary (35%)": 0.35, "Detailed Summary (55%)": 0.55}
    ratio = ratio_dict[length_ratio]
    
    summaries = []
    
    # We only show visual/download stats for the first doc if bulk uploaded
    for idx, doc_text in enumerate(docs):
        if not doc_text.strip():
            summaries.append("")
            continue
        try:
            if method == "Local Extractive (CPU & Fast)":
                sum_text = run_extraw_cpu(doc_text, ratio)
            else:
                sum_text = run_transformer_summarize(doc_text, hf_token, hf_model, ratio)
            summaries.append(sum_text)
        except Exception as e:
            return None, None, f"Execution failed at row {idx + 1}: {str(e)}"
            
    final_summary = summaries[0]
    original_len = len(docs[0].split())
    summary_len = len(final_summary.split())
    compression = round((1 - (summary_len / max(1, original_len))) * 100, 1)
    
    # Save output txt
    out_path = "document_summary.txt"
    with open(out_path, 'w', encoding='utf-8') as f:
        f.write(final_summary)
        
    # Clean visual metrics
    metrics_md = f"""
    ### Summarization Metrics
    - **Original Document Length**: {original_len} words
    - **Summary Length**: {summary_len} words
    - **Compression Rate**: {compression}% shorter than the original
    """
    
    return final_summary, out_path, metrics_md

custom_css = """
body {
    background-color: #0b0f19;
    color: #f3f4f6;
}
.gradio-container {
    font-family: 'Inter', sans-serif !important;
}
h1, h2 {
    color: #6366f1 !important;
}
"""

with gr.Blocks(theme=gr.themes.Default(primary_hue="indigo", secondary_hue="slate"), css=custom_css) as demo:
    df_state = gr.State()
    
    gr.HTML("""
    <div style="text-align: center; margin-bottom: 2rem;">
        <h1 style="font-size: 2.5rem; font-weight: 700; margin-bottom: 0.5rem; background: linear-gradient(to right, #6366f1, #a855f7); -webkit-background-clip: text; -webkit-text-fill-color: transparent;">Interactive Text Summarizer</h1>
        <p style="font-size: 1.1rem; color: #94a3b8; max-width: 800px; margin: 0 auto;">
            Condense long articles, book chapters, or reports down to essential summaries. 
            Runs locally on standard CPU scoring, or utilizes advanced neural models using your personal Hugging Face Token.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Choose Input Source")
            with gr.Tabs():
                with gr.TabItem("Paste Raw Text"):
                    text_input = gr.Textbox(
                        label="Source Text",
                        placeholder="Paste your document here (e.g., academic paper, book chapter, or news article)...",
                        lines=12
                    )
                with gr.TabItem("Upload Dataset File"):
                    file_input = gr.File(label="Upload (.csv, .xlsx, .txt)", file_types=[".csv", ".xlsx", ".txt"])
                    text_column_selector = gr.Dropdown(
                        label="Target Text Column", 
                        choices=[], 
                        visible=False,
                        interactive=True
                    )
                    status_text = gr.Markdown("No file uploaded yet.")
                    
            gr.Markdown("### 2. Configure Summarization")
            method_selector = gr.Radio(
                choices=["Local Extractive (CPU & Fast)", "Transformers (API Mode)"],
                value="Local Extractive (CPU & Fast)",
                label="Summarizer Model"
            )
            
            with gr.Group() as token_group:
                hf_token_input = gr.Textbox(
                    label="Hugging Face API Token",
                    placeholder="hf_...",
                    type="password",
                    visible=False,
                    info="Required to call advanced neural summarizers. Get one free at huggingface.co."
                )
                hf_model_input = gr.Dropdown(
                    choices=[
                        "facebook/bart-large-cnn",
                        "google/pegasus-xsum",
                        "philschmid/bart-large-cnn-samsum"
                    ],
                    value="facebook/bart-large-cnn",
                    label="Summarization Model (HF API)",
                    visible=False
                )
                
            length_selector = gr.Dropdown(
                choices=["Short Summary (15%)", "Medium Summary (35%)", "Detailed Summary (55%)"],
                value="Medium Summary (35%)",
                label="Target Summary Length"
            )
            
            run_btn = gr.Button("Generate Summary", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("### 3. Generated Summary Output")
            metrics_markdown = gr.Markdown("Summarization metrics will appear here after execution.")
            
            summary_output = gr.Textbox(
                label="Summary Content",
                lines=12,
                interactive=False
            )
            
            gr.Markdown("### 4. Export & Download")
            download_btn = gr.File(label="Download Summary Text File (.txt)")

    # Show/hide token field depending on model
    def toggle_method_fields(method):
        if method == "Transformers (API Mode)":
            return gr.update(visible=True), gr.update(visible=True)
        else:
            return gr.update(visible=False), gr.update(visible=False)
            
    method_selector.change(
        fn=toggle_method_fields,
        inputs=method_selector,
        outputs=[hf_token_input, hf_model_input]
    )
    
    file_input.change(
        fn=load_data,
        inputs=file_input,
        outputs=[df_state, text_column_selector, status_text]
    )
    
    run_btn.click(
        fn=process_summarization,
        inputs=[text_input, file_input, text_column_selector, method_selector, hf_token_input, hf_model_input, length_selector],
        outputs=[summary_output, download_btn, metrics_markdown]
    )

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