"""
🌙 DreamWeaver AI - Dream Journal & Interpreter (Local Transformers Version)
====================================================================
An innovative AI-powered dream analysis tool that interprets your dreams,
identifies emotional themes, and generates creative dream-inspired stories.

This version uses local Transformers pipelines (downloaded at runtime) for model execution.

Citations:
- Hugging Face Transformers Documentation
- Gradio Documentation (gradio.app)
- GitHub Copilot for code assistance
- Dream symbolism inspired by Jungian psychology concepts
"""

import gradio as gr
import time
import random
import os
from datetime import datetime


class _LabelScore:
    def __init__(self, label: str, score: float):
        self.label = label
        self.score = score


class LocalTransformersClient:
    """A tiny wrapper around local `transformers` pipelines.

    Keeps the same `client.text_classification()` and `client.text_generation()`
    interface used by the app and tests, but runs models locally.
    """

    def __init__(self):
        self._emotion_model_id: str | None = None
        self._emotion_pipe = None
        self._generation_model_id: str | None = None
        self._generation_pipe = None

    def _get_emotion_pipe(self, model: str):
        if self._emotion_pipe is not None and self._emotion_model_id == model:
            return self._emotion_pipe

        from transformers import pipeline

        self._emotion_model_id = model
        self._emotion_pipe = pipeline(
            task="text-classification",
            model=model,
            device=-1,
            top_k=None,
        )
        return self._emotion_pipe

    def _get_generation_pipe(self, model: str):
        if self._generation_pipe is not None and self._generation_model_id == model:
            return self._generation_pipe

        from transformers import pipeline

        self._generation_model_id = model
        self._generation_pipe = pipeline(
            task="text-generation",
            model=model,
            device=-1,
        )
        return self._generation_pipe

    def text_classification(self, text: str, model: str):
        pipe = self._get_emotion_pipe(model)
        raw = pipe(text)

        # transformers can return `[[{label,score},...]]` when running on a list of inputs.
        if isinstance(raw, list) and raw and isinstance(raw[0], list):
            raw = raw[0]

        if not isinstance(raw, list):
            raise RuntimeError(f"Unexpected text_classification response: {raw!r}")

        results: list[_LabelScore] = []
        for item in raw:
            if isinstance(item, dict) and "label" in item and "score" in item:
                results.append(_LabelScore(str(item["label"]), float(item["score"])))

        results.sort(key=lambda x: x.score, reverse=True)
        if not results:
            raise RuntimeError(f"Empty/invalid text_classification response: {raw!r}")
        return results

    def text_generation(
        self,
        prompt: str,
        model: str,
        max_new_tokens: int = 200,
        temperature: float = 0.7,
        do_sample: bool = True,
    ) -> str:
        pipe = self._get_generation_pipe(model)
        try:
            out = pipe(
                prompt,
                max_new_tokens=int(max_new_tokens),
                temperature=float(temperature),
                do_sample=bool(do_sample),
                return_full_text=False,
            )
        except TypeError:
            out = pipe(
                prompt,
                max_new_tokens=int(max_new_tokens),
                temperature=float(temperature),
                do_sample=bool(do_sample),
            )

        if not isinstance(out, list) or not out or not isinstance(out[0], dict) or "generated_text" not in out[0]:
            raise RuntimeError(f"Unexpected text_generation response: {out!r}")

        generated = str(out[0]["generated_text"])
        if generated.startswith(prompt):
            generated = generated[len(prompt):].lstrip()
        return generated


client = LocalTransformersClient()

# Dream symbol database for enhanced interpretations
DREAM_SYMBOLS = {
    "water": "💧 Emotions, unconscious mind, purification",
    "flying": "🦋 Freedom, ambition, escaping limitations", 
    "falling": "⬇️ Loss of control, anxiety, letting go",
    "teeth": "🦷 Confidence, self-image, communication",
    "chase": "🏃 Avoidance, pressure, unresolved issues",
    "house": "🏠 Self, psyche, different aspects of personality",
    "snake": "🐍 Transformation, hidden fears, healing",
    "death": "💀 Endings, transformation, new beginnings",
    "baby": "👶 New beginnings, innocence, vulnerability",
    "car": "🚗 Life direction, control, personal drive",
    "fire": "🔥 Passion, anger, transformation, energy",
    "ocean": "🌊 Vast emotions, the unknown, life's depth",
    "forest": "🌲 Unconscious, mystery, personal growth",
    "mirror": "🪞 Self-reflection, truth, identity",
    "stairs": "🪜 Progress, transition, spiritual journey",
    "door": "🚪 Opportunities, transitions, new paths",
    "rain": "🌧️ Cleansing, sadness, renewal",
    "sun": "☀️ Clarity, vitality, consciousness",
    "moon": "🌙 Intuition, feminine energy, cycles",
    "bird": "🐦 Freedom, perspective, spiritual messages"
}

# Mood colors for visualization
MOOD_COLORS = {
    "joy": "#FFD700",
    "fear": "#4B0082", 
    "sadness": "#4169E1",
    "anger": "#DC143C",
    "surprise": "#FF69B4",
    "peace": "#90EE90",
    "confusion": "#DDA0DD",
    "excitement": "#FF4500"
}


def analyze_dream_sentiment(dream_text: str) -> tuple:
    """
    Analyze the emotional content of a dream using sentiment analysis.
    """
    if not dream_text.strip():
        return "Please describe your dream first.", "N/A", ""
    
    start_time = time.time()
    
    try:
        # Multi-label emotion classification
        emotions = client.text_classification(
            dream_text,
            model="j-hartmann/emotion-english-distilroberta-base"
        )
        
        end_time = time.time()
        response_time = f"{(end_time - start_time):.3f}s"
        
        # Format emotional analysis
        result = "## 🎭 Emotional Landscape of Your Dream\n\n"
        
        emotion_bars = ""
        for emotion in emotions[:5]:  # Top 5 emotions
            label = emotion.label.capitalize()
            score = emotion.score
            bar_length = int(score * 20)
            color = MOOD_COLORS.get(label.lower(), "#888888")
            emoji = {"joy": "😊", "fear": "😨", "sadness": "😢", "anger": "😠", 
                    "surprise": "😲", "disgust": "🤢", "neutral": "😐"}.get(label.lower(), "🔮")
            
            result += f"{emoji} **{label}**: {'█' * bar_length}{'░' * (20-bar_length)} {score:.1%}\n"
        
        # Dominant mood
        dominant = emotions[0].label if emotions else "Unknown"
        result += f"\n### 🎯 Dominant Mood: **{dominant.upper()}**"
        
        return result, response_time, dominant.lower()
        
    except Exception as e:
        return f"❌ Error analyzing emotions: {str(e)}", "N/A", ""


def generate_text(prompt: str, max_tokens: int, temperature: float) -> tuple:
    """Basic text generation helper (used by unit tests)."""
    if not prompt.strip():
        return "Please enter some text to generate.", "N/A"

    start_time = time.time()
    try:
        result = client.text_generation(
            prompt,
                model="distilgpt2",
            max_new_tokens=max_tokens,
            temperature=temperature,
            do_sample=True,
        )
        end_time = time.time()
        response_time = f"{(end_time - start_time):.3f}s"
        return result, response_time
    except Exception as e:
        return f"❌ Error generating text: {str(e)}", "N/A"


def find_dream_symbols(dream_text: str) -> str:
    """
    Identify and interpret common dream symbols in the text.
    """
    dream_lower = dream_text.lower()
    found_symbols = []
    
    for symbol, meaning in DREAM_SYMBOLS.items():
        if symbol in dream_lower:
            found_symbols.append(f"- **{symbol.capitalize()}**: {meaning}")
    
    if found_symbols:
        return "## 🔮 Dream Symbols Detected\n\n" + "\n".join(found_symbols)
    else:
        return "## 🔮 Dream Symbols\n\n*No common symbols detected. Your dream may contain unique personal symbolism.*"


def generate_dream_interpretation(dream_text: str, mood: str) -> tuple:
    """
    Generate an AI interpretation of the dream using text generation.
    """
    if not dream_text.strip():
        return "Please describe your dream first.", "N/A"
    
    start_time = time.time()
    
    try:
        # Create a prompt for dream interpretation
        prompt = f"""As a dream analyst, provide a thoughtful interpretation of this dream:

Dream: {dream_text}

The dreamer's dominant emotion was: {mood}

Interpretation:"""

        result = client.text_generation(
            prompt,
                model="distilgpt2",
            max_new_tokens=200,
            temperature=0.7,
            do_sample=True
        )
        
        end_time = time.time()
        response_time = f"{(end_time - start_time):.3f}s"
        
        interpretation = f"## 📖 Dream Interpretation\n\n{result}"
        return interpretation, response_time
        
    except Exception as e:
        return f"❌ Error generating interpretation: {str(e)}", "N/A"


def generate_dream_story(dream_text: str, genre: str, length: int) -> tuple:
    """
    Transform the dream into a creative short story.
    """
    if not dream_text.strip():
        return "Please describe your dream first.", "N/A"
    
    start_time = time.time()
    
    try:
        genre_prompts = {
            "Fantasy": "Write a magical fantasy story",
            "Sci-Fi": "Write a futuristic science fiction story", 
            "Mystery": "Write a mysterious thriller story",
            "Romance": "Write a romantic story",
            "Horror": "Write a suspenseful horror story",
            "Adventure": "Write an exciting adventure story"
        }
        
        prompt = f"""{genre_prompts.get(genre, "Write a creative story")} inspired by this dream:

Dream elements: {dream_text}

Story:"""

        result = client.text_generation(
            prompt,
                model="distilgpt2",
            max_new_tokens=length,
            temperature=0.8,
            do_sample=True
        )
        
        end_time = time.time()
        response_time = f"{(end_time - start_time):.3f}s"
        
        story = f"## ✨ Dream-Inspired {genre} Story\n\n{result}"
        return story, response_time
        
    except Exception as e:
        return f"❌ Error generating story: {str(e)}", "N/A"


def generate_dream_image_prompt(dream_text: str) -> str:
    """
    Generate an image prompt based on the dream for use with image generators.
    """
    # Extract key visual elements
    visual_keywords = []
    
    for symbol in DREAM_SYMBOLS.keys():
        if symbol in dream_text.lower():
            visual_keywords.append(symbol)
    
    # Add atmospheric descriptors based on common dream themes
    atmosphere = random.choice([
        "surreal", "ethereal", "mystical", "dreamlike", 
        "fantastical", "otherworldly", "atmospheric"
    ])
    
    style = random.choice([
        "digital art", "oil painting style", "watercolor", 
        "concept art", "impressionist", "fantasy illustration"
    ])
    
    if visual_keywords:
        elements = ", ".join(visual_keywords[:5])
        prompt = f"A {atmosphere} scene featuring {elements}, {style}, dreamy lighting, vivid colors, detailed"
    else:
        prompt = f"A {atmosphere} dreamscape, {style}, surreal environment, dreamy lighting, mysterious atmosphere"
    
    return f"## 🎨 Image Generation Prompt\n\n*Use this prompt with DALL-E, Midjourney, or Stable Diffusion:*\n\n```\n{prompt}\n```"


def save_to_journal(dream_text: str, interpretation: str) -> str:
    """
    Format dream entry for saving to a journal.
    """
    timestamp = datetime.now().strftime("%B %d, %Y at %I:%M %p")
    
    journal_entry = f"""
# 📔 Dream Journal Entry
## {timestamp}

### 💭 The Dream
{dream_text}

### 🔮 Interpretation
{interpretation}

---
*Recorded with DreamWeaver AI*
"""
    return journal_entry


# Create the main Gradio interface
with gr.Blocks(
    title="🌙 DreamWeaver AI",
    theme=gr.themes.Soft(
        primary_hue="indigo",
        secondary_hue="purple",
        neutral_hue="slate"
    ),
    css="""
    .gradio-container {
        background: linear-gradient(135deg, #1a1a2e 0%, #16213e 50%, #0f3460 100%);
    }
    .main-title {
        text-align: center;
        color: #e2e2e2;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.5);
    }
    """
) as demo:
    
    gr.Markdown("""
    # 🌙 DreamWeaver AI - Dream Journal & Interpreter
    ### *Unlock the secrets of your subconscious mind*
    
    Welcome to DreamWeaver AI! This innovative tool uses advanced AI to analyze your dreams,
    identify emotional patterns, discover symbolic meanings, and even transform your dreams
    into creative stories.
    
    **✨ Local Transformers Version** - Models are downloaded and run on this Space
    
    ---
    """)
    
    with gr.Row():
        with gr.Column(scale=2):
            dream_input = gr.Textbox(
                label="🌙 Describe Your Dream",
                placeholder="Last night I dreamed I was flying over a vast ocean. The water was crystal clear and I could see colorful fish below. Suddenly, I noticed a mysterious door floating in the sky...",
                lines=8,
                max_lines=15
            )
            
            with gr.Row():
                analyze_btn = gr.Button("🔍 Analyze Dream", variant="primary", size="lg")
                clear_btn = gr.Button("🗑️ Clear", variant="secondary")
    
    with gr.Tabs():
        # Tab 1: Emotional Analysis
        with gr.TabItem("🎭 Emotional Analysis"):
            with gr.Row():
                with gr.Column():
                    emotion_output = gr.Markdown(label="Emotional Landscape")
                with gr.Column():
                    emotion_time = gr.Textbox(label="⏱️ Analysis Time", interactive=False)
            
            mood_state = gr.State("")
        
        # Tab 2: Symbol Interpretation
        with gr.TabItem("🔮 Dream Symbols"):
            symbols_output = gr.Markdown(label="Detected Symbols")
        
        # Tab 3: AI Interpretation
        with gr.TabItem("📖 AI Interpretation"):
            with gr.Row():
                interpret_btn = gr.Button("🧠 Generate Interpretation", variant="primary")
            interpretation_output = gr.Markdown(label="Dream Interpretation")
            interpretation_time = gr.Textbox(label="⏱️ Generation Time", interactive=False)
        
        # Tab 4: Dream Story Generator
        with gr.TabItem("✨ Story Generator"):
            gr.Markdown("### Transform your dream into a creative story!")
            
            with gr.Row():
                genre_dropdown = gr.Dropdown(
                    choices=["Fantasy", "Sci-Fi", "Mystery", "Romance", "Horror", "Adventure"],
                    value="Fantasy",
                    label="📚 Select Genre"
                )
                length_slider = gr.Slider(
                    minimum=100,
                    maximum=500,
                    value=250,
                    step=50,
                    label="📏 Story Length (tokens)"
                )
            
            story_btn = gr.Button("✍️ Generate Story", variant="primary")
            story_output = gr.Markdown(label="Your Dream Story")
            story_time = gr.Textbox(label="⏱️ Generation Time", interactive=False)
        
        # Tab 5: Image Prompt
        with gr.TabItem("🎨 Visualize"):
            gr.Markdown("### Create visual art from your dream!")
            image_prompt_btn = gr.Button("🖼️ Generate Image Prompt", variant="primary")
            image_prompt_output = gr.Markdown(label="Image Prompt")
        
        # Tab 6: Dream Journal
        with gr.TabItem("📔 Journal"):
            gr.Markdown("### Save your dream to your journal")
            save_btn = gr.Button("💾 Format for Journal", variant="primary")
            journal_output = gr.Textbox(
                label="Journal Entry (Copy this)",
                lines=15,
                show_copy_button=True
            )
        
        # Tab 7: About
        with gr.TabItem("ℹ️ About"):
            gr.Markdown("""
            ## About DreamWeaver AI
            
            ### 🌟 Features
            - **Emotional Analysis**: Detect the emotional undertones of your dreams
            - **Symbol Detection**: Identify and interpret common dream symbols
            - **AI Interpretation**: Get personalized dream interpretations
            - **Story Generation**: Transform dreams into creative stories
            - **Visual Prompts**: Generate prompts for AI image generators
            
            ### 🔧 Technical Architecture
            
            | Component | Technology |
            |-----------|------------|
            | Frontend | Gradio 4.x |
            | Emotion Model | distilroberta-base |
            | Text Generation | distilgpt2 |
            | Hosting | Hugging Face Spaces |
            | Inference | Local Transformers |
            
            ### ⚡ Local Model Approach Trade-offs
            
            | ✅ Advantages | ⚠️ Considerations |
            |--------------|-------------------|
            | No external API dependency | First-run model download time |
            | No network latency | More RAM/disk usage |
            | Works offline (after download) | Smaller models on CPU |
            
            ### 📚 Citations
            - Hugging Face Transformers
            - Gradio Documentation
            - Dream symbolism: Jungian psychology concepts
            - GitHub Copilot assistance
            
            ### 👥 Team
            *Add your team members here*
            """)
    
    # Example dreams
    gr.Examples(
        examples=[
            ["I was flying over a beautiful forest at sunset. The trees below were golden and I felt completely free. Suddenly I noticed I was being chased by a dark shadow."],
            ["I found myself in my childhood home, but all the rooms were different. There was a mysterious door that I had never seen before. When I opened it, I saw an endless staircase."],
            ["I was swimming in a crystal-clear ocean with colorful fish. The water was warm and I could breathe underwater. I discovered an ancient underwater city with golden buildings."],
            ["I was taking an important exam but realized I couldn't read any of the questions. My teeth started falling out and everyone was staring at me."],
        ],
        inputs=dream_input,
        label="💭 Example Dreams"
    )
    
    # Event handlers
    def full_analysis(dream_text):
        emotions, time, mood = analyze_dream_sentiment(dream_text)
        symbols = find_dream_symbols(dream_text)
        return emotions, time, mood, symbols
    
    analyze_btn.click(
        fn=full_analysis,
        inputs=[dream_input],
        outputs=[emotion_output, emotion_time, mood_state, symbols_output]
    )
    
    interpret_btn.click(
        fn=generate_dream_interpretation,
        inputs=[dream_input, mood_state],
        outputs=[interpretation_output, interpretation_time]
    )
    
    story_btn.click(
        fn=generate_dream_story,
        inputs=[dream_input, genre_dropdown, length_slider],
        outputs=[story_output, story_time]
    )
    
    image_prompt_btn.click(
        fn=generate_dream_image_prompt,
        inputs=[dream_input],
        outputs=[image_prompt_output]
    )
    
    save_btn.click(
        fn=lambda d, i: save_to_journal(d, i),
        inputs=[dream_input, interpretation_output],
        outputs=[journal_output]
    )
    
    clear_btn.click(
        fn=lambda: ("", "", "", "", "", "", "", ""),
        outputs=[dream_input, emotion_output, emotion_time, symbols_output, 
                interpretation_output, story_output, image_prompt_output, journal_output]
    )
    
    gr.Markdown("""
    ---
    <center>
    
    *🌙 DreamWeaver AI - Powered by Transformers*
    
    *Made with ❤️ for MLOps Case Study*
    
    </center>
    """)


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