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)}"

# Local lexical mock style dictionary for fun offline demonstrations
STYLE_REPLACEMENTS = {
    "Shakespearean (Early Modern)": {
        r'\byou\b': "thou",
        r'\bare\b': "art",
        r'\bdo\b': "dost",
        r'\bdoes\b': "doth",
        r'\bhave\b': "hast",
        r'\bhas\b': "hath",
        r'\bknow\b': "wot",
        r'\bplease\b': "pray",
        r'\bwrite\b': "indite",
        r'\bwords\b': "runes",
        r'\bfriend\b': "comrade",
        r'\bhappy\b': "merry",
        r'\bbefore\b': "ere",
        r'\bwhy\b': "wherefore",
        r'\bhello\b': "hark",
        r'\bgoodbye\b': "fare thee well"
    },
    "Orwellian (Dystopian Bureaucracy)": {
        r'\bgood\b': "plusgood",
        r'\bexcellent\b': "doubleplusgood",
        r'\bbad\b': "ungood",
        r'\btruth\b': "minitrue",
        r'\bpeace\b': "minipax",
        r'\bthought\b': "crimethink",
        r'\bwork\b': "labor-core",
        r'\bfreedom\b': "slavery",
        r'\bhistory\b': "rectypo",
        r'\bthink\b': "doublethink"
    },
    "Hemingwayesque (Minimalist)": {
        # Hemingway is about removing long adjectives/adverbs, we mock it locally
        r'\bextremely\b': "",
        r'\bvery\b': "",
        r'\bbeautifully\b': "",
        r'\binterestingly\b': "",
        r'\bsuddenly\b': "",
        r'\bcompletely\b': ""
    }
}

def run_local_style(text, style):
    """Offline rule-based lexical style mapper."""
    styled_text = text
    
    # Apply standard word replacements
    if style in STYLE_REPLACEMENTS:
        replacements = STYLE_REPLACEMENTS[style]
        for pattern, replacement in replacements.items():
            # Match case-insensitively but retain capitalization
            def repl_func(match):
                word = match.group()
                if word.istitle():
                    return replacement.capitalize()
                elif word.isupper():
                    return replacement.upper()
                return replacement
                
            styled_text = re.sub(pattern, repl_func, styled_text, flags=re.IGNORECASE)
            
    # Add a stylistic suffix/flavor to ensure immediate fun results
    if style == "Shakespearean (Early Modern)":
        styled_text = f"Hark! {styled_text}—Anon, thy runes have spoken!"
    elif style == "Orwellian (Dystopian Bureaucracy)":
        styled_text = f"{styled_text} Big Brother approves this message."
    elif style == "Hemingwayesque (Minimalist)":
        # clean extra spaces from deleted adverbs
        styled_text = re.sub(r'\s+', ' ', styled_text).strip()
        styled_text = f"{styled_text} It was good. The sun was hot."
    elif style == "Academic / Formal":
        styled_text = f"It is widely postulated that: {styled_text} Hence, subsequent empirical inquiries are mandated."
        
    return styled_text

def run_neural_style(text, hf_token, model_name, style):
    """Uses generative models to fully rewrite draft text into target style."""
    if not hf_token:
        raise ValueError("Hugging Face API Access Token is required for Transformers mode.")
        
    client = InferenceClient(token=hf_token)
    
    prompt = f"""[INST] Rewrite this text exactly in the writing style of {style}. Preserve the core meaning but transform the sentence structure, tone, word choice, and cadence to sound highly authentic to the target style.
Do not output extra text, commentary, or markdown tags. Just output the rewritten text.

Text to rewrite:
"{text}" [/INST]"""
    
    try:
        response = client.text_generation(prompt, model=model_name, max_new_tokens=400, temperature=0.6)
        return response.strip()
    except Exception as e:
        raise RuntimeError(f"Hugging Face API error: {str(e)}")

def process_style_transfer(text_input, file_obj, text_col, method, hf_token, hf_model, target_style):
    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."
        
    try:
        if method == "Local Lexical Replacer (CPU & Fast)":
            styled_text = run_local_style(docs[0], target_style)
        else:
            styled_text = run_neural_style(docs[0], hf_token, hf_model, target_style)
            
        # Save output txt
        out_path = "styled_document.txt"
        with open(out_path, 'w', encoding='utf-8') as f:
            f.write(styled_text)
            
        status_md = f"Style transfer complete: Transformed first document into **{target_style}** style."
        
        return styled_text, out_path, status_md
        
    except Exception as e:
        return None, None, f"Execution failed: {str(e)}"

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;">Writing Style Transfer</h1>
        <p style="font-size: 1.1rem; color: #94a3b8; max-width: 800px; margin: 0 auto;">
            Apply the literary writing style of iconic authors—Shakespeare, Orwell, or Hemingway—to your own text. 
            Experiment with local keyword replacements or unlock advanced neural style transfers using your Hugging Face Token.
        </p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 1. Source Document")
            with gr.Tabs():
                with gr.TabItem("Paste Raw Text"):
                    text_input = gr.Textbox(
                        label="Source Text",
                        placeholder="Paste draft text here (e.g., 'You should write a message to your friend and tell them to have a nice day.')...",
                        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 Transfer")
            method_selector = gr.Radio(
                choices=["Local Lexical Replacer (CPU & Fast)", "Transformers (API Mode)"],
                value="Local Lexical Replacer (CPU & Fast)",
                label="Transfer 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 LLMs. Get one free at huggingface.co."
                )
                hf_model_input = gr.Dropdown(
                    choices=[
                        "Qwen/Qwen2.5-7B-Instruct",
                        "meta-llama/Llama-3-8b-instruct",
                        "mistralai/Mistral-7B-Instruct-v0.3"
                    ],
                    value="Qwen/Qwen2.5-7B-Instruct",
                    label="Transformer Model (HF API)",
                    visible=False
                )
                
            target_style = gr.Dropdown(
                choices=["Shakespearean (Early Modern)", "Orwellian (Dystopian Bureaucracy)", "Hemingwayesque (Minimalist)", "Academic / Formal"],
                value="Shakespearean (Early Modern)",
                label="Target Writing Style"
            )
            
            run_btn = gr.Button("Transform Style", variant="primary")
            
        with gr.Column(scale=2):
            gr.Markdown("### 3. Styled Document Output")
            status_markdown = gr.Markdown("Enter draft text and click 'Transform Style' to run.")
            
            styled_output = gr.Textbox(
                label="Styled Document",
                lines=15,
                interactive=False
            )
            
            gr.Markdown("### 4. Export")
            download_btn = gr.File(label="Download Styled 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_style_transfer,
        inputs=[text_input, file_input, text_column_selector, method_selector, hf_token_input, hf_model_input, target_style],
        outputs=[styled_output, download_btn, status_markdown]
    )

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