"""
Hugging Face Gradio App for Iain Morris Style Article Generation
"""

import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from peft import PeftModel
import logging
import os

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class IainMorrisGenerator:
    def __init__(self):
        """Initialize the article generator"""
        self.base_model_name = "HuggingFaceH4/zephyr-7b-beta"
        self.lora_adapter_path = "models/iain-morris-model-enhanced"
        
        # Check for Apple Silicon MPS support
        if torch.backends.mps.is_available():
            self.device = torch.device("mps")
        elif torch.cuda.is_available():
            self.device = torch.device("cuda")
        else:
            self.device = torch.device("cpu")
        
        # Quantization config for inference (only for CUDA)
        if self.device.type == "cuda":
            self.bnb_config = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_use_double_quant=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_compute_dtype=torch.bfloat16
            )
        else:
            self.bnb_config = None
        
        self.model = None
        self.tokenizer = None
        self.system_prompt = """You are Iain Morris, the cynical telecom journalist from Light Reading with a gift for doom-laden analysis and visceral metaphors. Your writing style is unmistakable:

**PROVOCATIVE DOOM-LADEN OPENINGS**: Start with impending disaster, catastrophe, or failure
- "What could possibly go wrong?" (signature phrase)
- "train wreck", "disaster", "catastrophe", "nightmare scenario"
- Present optimistic industry claims, then systematically demolish them

**SIGNATURE DARK ANALOGIES**: Use visceral, physical metaphors for abstract concepts
- Technology as disease, infection, or bodily harm
- Business as warfare, natural disasters, or medical procedures
- Markets as living organisms that can be sick, dying, or mutating

**CYNICAL WIT & EXPERTISE**: Combine deep technical knowledge with cutting observations
- Parenthetical snark (like this, but funnier and more cutting)
- Quote industry executives, then undercut them immediately
- Show technical expertise while mocking industry groupthink

**DISTINCTIVE PHRASES**: Weave in signature expressions
- "What could possibly go wrong?"
- "train wreck" / "disaster" / "catastrophe"
- "Meanwhile, back in reality..."
- "Of course, there's a catch" / "But here's the thing"
- References to "drinking the Kool-Aid" or similar

**BRITISH CYNICISM**: Dry, cutting observations about human nature and corporate behavior
- Assume the worst about corporate motives
- Highlight contradictions and hypocrisy
- Use understatement for dramatic effect

Write a compelling article that captures this distinctive voice - cynical, expert, and darkly entertaining."""

    def load_model(self):
        """Load the fine-tuned model"""
        try:
            logger.info("Loading tokenizer...")
            self.tokenizer = AutoTokenizer.from_pretrained(
                self.base_model_name,
                trust_remote_code=True,
                padding_side="left"
            )
            
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
            
            logger.info(f"Loading base model on device: {self.device}")
            
            # Configure model loading based on device
            model_kwargs = {
                "trust_remote_code": True,
                "torch_dtype": torch.bfloat16
            }
            
            if self.device.type == "cuda":
                model_kwargs["quantization_config"] = self.bnb_config
                model_kwargs["device_map"] = "auto"
            else:
                # For MPS or CPU, load without quantization
                model_kwargs["device_map"] = None
            
            base_model = AutoModelForCausalLM.from_pretrained(
                self.base_model_name,
                **model_kwargs
            )
            
            # Move to device if not using device_map
            if model_kwargs["device_map"] is None:
                base_model = base_model.to(self.device)
            
            # Load LoRA adapters if they exist
            if os.path.exists(self.lora_adapter_path):
                logger.info("Loading LoRA adapters...")
                self.model = PeftModel.from_pretrained(base_model, self.lora_adapter_path)
            else:
                logger.warning("LoRA adapters not found. Using base model.")
                self.model = base_model
            
            logger.info("Model loaded successfully!")
            return True
            
        except Exception as e:
            logger.error(f"Error loading model: {e}")
            return False

    def generate_article(self, topic: str, max_length: int = 1000, temperature: float = 0.7, top_p: float = 0.9) -> str:
        """
        Generate an article on the given topic
        
        Args:
            topic: Topic to write about
            max_length: Maximum length of generated text
            temperature: Sampling temperature
            top_p: Top-p sampling parameter
            
        Returns:
            Generated article text
        """
        if self.model is None or self.tokenizer is None:
            return "Error: Model not loaded. Please wait for the model to initialize."
        
        try:
            # Create the prompt
            messages = [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": f"Write a telecom industry news article about: {topic}"}
            ]
            
            # Format the prompt
            if hasattr(self.tokenizer, 'apply_chat_template'):
                try:
                    prompt = self.tokenizer.apply_chat_template(
                        messages, 
                        tokenize=False, 
                        add_generation_prompt=True
                    )
                except:
                    # Fallback formatting
                    prompt = f"<|system|>\n{self.system_prompt}\n<|user|>\nWrite a telecom industry news article about: {topic}\n<|assistant|>\n"
            else:
                prompt = f"<|system|>\n{self.system_prompt}\n<|user|>\nWrite a telecom industry news article about: {topic}\n<|assistant|>\n"
            
            # Tokenize
            inputs = self.tokenizer(
                prompt,
                return_tensors="pt",
                truncation=True,
                max_length=1024
            ).to(self.device)
            
            # Generate
            with torch.no_grad():
                outputs = self.model.generate(
                    **inputs,
                    max_new_tokens=max_length,
                    temperature=temperature,
                    top_p=top_p,
                    do_sample=True,
                    pad_token_id=self.tokenizer.eos_token_id,
                    repetition_penalty=1.1,
                    no_repeat_ngram_size=3
                )
            
            # Decode
            generated_text = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            
            # Extract only the generated part (after the prompt)
            if "<|assistant|>" in generated_text:
                article = generated_text.split("<|assistant|>")[-1].strip()
            else:
                # Fallback: remove the original prompt
                article = generated_text[len(prompt):].strip()
            
            return article
            
        except Exception as e:
            logger.error(f"Error generating article: {e}")
            return f"Error generating article: {str(e)}"

# Initialize the generator
generator = IainMorrisGenerator()

def generate_article_interface(topic, max_length, temperature, top_p):
    """Interface function for Gradio"""
    if not topic.strip():
        return "Please enter a topic for the article."
    
    return generator.generate_article(topic, max_length, temperature, top_p)

def load_model_interface():
    """Load model interface for Gradio"""
    success = generator.load_model()
    if success:
        return "✅ Model loaded successfully! You can now generate articles."
    else:
        return "❌ Failed to load model. Please check the logs."

# Create Gradio interface
with gr.Blocks(
    title="Iain Morris Style Telecom Article Generator",
    theme=gr.themes.Soft(),
    css="""
    .main-header {
        text-align: center;
        background: linear-gradient(90deg, #1e3a8a, #3b82f6);
        color: white;
        padding: 2rem;
        border-radius: 10px;
        margin-bottom: 2rem;
    }
    .description {
        text-align: center;
        font-size: 1.1em;
        color: #4b5563;
        margin-bottom: 2rem;
    }
    """
) as demo:
    
    gr.HTML("""
    <div class="main-header">
        <h1 style="color: white;">🗞️ The Morris-Bot</h1>
        <p style="color: white;">Generate telecom industry news articles in the distinctive style of Iain Morris from Light Reading</p>
    </div>
    """)
    
    gr.HTML("""
    <div class="description">
        <p>This AI model has been fine-tuned on articles by Iain Morris to capture his analytical writing style, 
        technical expertise, and slightly irreverent tone. Enter a telecom topic below to generate an article 
        in his distinctive voice.</p>
        <p><strong>⚠️ Please Note:</strong> Article generation may take several minutes due to the use of Hugging Face's free CPU infrastructure. Thank you for your patience!</p>
    </div>
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            # Model loading section
            gr.Markdown("### 🚀 Model Setup")
            load_btn = gr.Button("Load Model", variant="primary", size="lg")
            load_status = gr.Textbox(
                label="Status",
                value="Click 'Load Model' to initialize the AI model",
                interactive=False
            )
            
            gr.Markdown("### 📝 Article Generation")
            topic_input = gr.Textbox(
                label="Article Topic",
                placeholder="e.g., 5G network slicing challenges for operators",
                lines=2
            )
            
            with gr.Accordion("Advanced Settings", open=False):
                max_length = gr.Slider(
                    minimum=200,
                    maximum=2000,
                    value=800,
                    step=50,
                    label="Max Article Length (tokens)"
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.5,
                    value=0.7,
                    step=0.1,
                    label="Creativity (Temperature)"
                )
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.9,
                    step=0.05,
                    label="Focus (Top-p)"
                )
            
            generate_btn = gr.Button("Generate Article", variant="secondary", size="lg")
        
        with gr.Column(scale=2):
            gr.Markdown("### 📰 Generated Article")
            output = gr.Textbox(
                label="Article Output",
                lines=25,
                placeholder="Generated article will appear here...",
                show_copy_button=True
            )
    
    # Example topics
    gr.Markdown("### 💡 Example Topics")
    examples = gr.Examples(
        examples=[
            ["OpenRAN deployment challenges for mobile operators"],
            ["The impact of AI on network automation"],
            ["Fiber vs 5G: The broadband infrastructure debate"],
            ["Edge computing adoption in telecommunications"],
            ["Regulatory challenges for satellite internet providers"],
            ["The future of network slicing in enterprise 5G"],
        ],
        inputs=[topic_input],
        label="Click an example to try it out"
    )
    
    # Footer
    gr.HTML("""
    <div style="text-align: center; margin-top: 2rem; padding: 1rem; background-color: #f9fafb; border-radius: 8px;">
        <p><strong>About:</strong> This model generates articles in the style of Iain Morris, 
        a respected telecom journalist known for his insightful analysis and accessible explanations 
        of complex technical topics.</p>
        <p><em>Generated content is AI-created and should be reviewed before publication.</em></p>
    </div>
    """)
    
    # Event handlers
    load_btn.click(
        fn=load_model_interface,
        outputs=[load_status]
    )
    
    generate_btn.click(
        fn=generate_article_interface,
        inputs=[topic_input, max_length, temperature, top_p],
        outputs=[output]
    )
    
    # Auto-generate on Enter key
    topic_input.submit(
        fn=generate_article_interface,
        inputs=[topic_input, max_length, temperature, top_p],
        outputs=[output]
    )

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )