import gradio as gr
import xml.etree.ElementTree as ET
import re
import anthropic
import os
from dotenv import load_dotenv
from utils import svg_to_png_base64
import pathlib
from logger import setup_logger

load_dotenv(override=True)
logger = setup_logger()


class SVGAnimationGenerator:
    def __init__(self):
        self.client = None
        self.predict_decompose_group_prompt = self._get_prompt(
            "prompts/predict_decompose_group.txt"
        )
        self.feedback_decompose_group_prompt = self._get_prompt(
            "prompts/feedback_decompose_group.txt"
        )
        self.generate_animation_prompt = self._get_prompt(
            "prompts/generate_animation.txt"
        )
        if "ANTHROPIC_API_KEY" in os.environ:
            self.client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

    def _get_prompt(self, prompt_file_path: str) -> str:
        try:
            with open(prompt_file_path, "r", encoding="utf-8") as f:
                return f.read()
        except FileNotFoundError:
            return "Prompt file not found. Please check the path."

    def parse_svg(self, svg_content: str) -> dict:
        try:
            # Remove namespace definitions and simplify the <svg> tag
            svg_content = re.sub(r'xmlns[^=]*="[^"]*"', "", svg_content)
            svg_content = re.sub(r"<svg[^>]*>", "<svg>", svg_content)
            return {"svg_content": svg_content}
        except Exception as e:
            return {"error": f"SVG parsing error: {e}"}

    def predict_decompose_group(self, parsed_svg: dict, object_name: str) -> dict:
        try:
            svg_content = parsed_svg["svg_content"]

            # Convert SVG to PNG for MLLM analysis
            image_media_type, image_data = svg_to_png_base64(svg_content)

            if not image_data:
                return {"error": "Failed to convert SVG to PNG"}

            prompt = self.predict_decompose_group_prompt.format(
                object_name=object_name, svg_content=svg_content
            )
            logger.info(f"Decomposition Prompt for {object_name}:\n{prompt}")

            response = self.client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=10000,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": image_media_type,
                                    "data": image_data,
                                },
                            },
                            {"type": "text", "text": prompt},
                        ],
                    },
                    {
                        "role": "assistant",
                        "content": [
                            {
                                "type": "text",
                                "text": "<animation_plan>"
                            }
                        ]
                    }
                ],
            )
            response_text = response.content[0].text
            logger.info(f"Model Response:\n{response_text}")

            decomposed_svg_match = re.search(
                r"<decomposed_svg>(.*?)</decomposed_svg>", response_text, re.DOTALL
            )
            animation_suggenstions_match = re.search(
                r"<animation_suggestions>(.*?)</animation_suggestions>",
                response_text,
                re.DOTALL,
            )
            print(
                "[SVG Decompose] Decomposed SVG found", decomposed_svg_match is not None
            )
            if decomposed_svg_match and animation_suggenstions_match:
                decomposed_svg_text = decomposed_svg_match.group(1).strip()
                animation_suggestions = animation_suggenstions_match.group(1).strip()
                print("[SVG Decompose] Animation suggestions found")
                return {
                    "svg_content": decomposed_svg_text,
                    "animation_suggestions": animation_suggestions,
                }
            else:
                return {
                    "error": "Decomposed SVG and Animation Suggestion not found in response."
                }
        except Exception as e:
            return {"error": f"Error during MLLM prediction: {e}"}

    def feedback_decompose_group(self, svg_content: str, feedback: str) -> tuple:
        try:
            # Parse the SVG content first
            parsed_svg = self.parse_svg(svg_content)  # Remove ["svg_content"] access
            if "error" in parsed_svg:
                error_message = parsed_svg["error"]
                error_html = create_error_html(error_message)
                return "", "", error_html  # Return tuple of 3 values

            prompt = self.feedback_decompose_group_prompt.format(
                parsed_svg=parsed_svg, feedback=feedback
            )
            logger.info(f"Feedback Prompt:\n{prompt}")

            response = self.client.messages.create(
                model="claude-sonnet-4-20250514",
                max_tokens=10000,
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": prompt},
                        ],
                    }
                ],
            )
            response_text = response.content[0].text
            logger.info(f"Model Response:\n{response_text}")

            decomposed_svg_match = re.search(
                r"<decomposed_svg>(.*?)</decomposed_svg>", response_text, re.DOTALL
            )
            animation_suggenstions_match = re.search(
                r"<animation_suggestions>(.*?)</animation_suggestions>",
                response_text,
                re.DOTALL,
            )

            if decomposed_svg_match and animation_suggenstions_match:
                decomposed_svg_text = decomposed_svg_match.group(1).strip()
                animation_suggestions = animation_suggenstions_match.group(1).strip()
                
                # Create viewer HTML
                viewer_html = f"""
                <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
                    <div style='display: block; align-items: center; margin-bottom: 10px;'>
                    </div>
                    <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
                        {decomposed_svg_text}
                    </div>
                </div>
                """
                
                return decomposed_svg_text, animation_suggestions, viewer_html  # Return tuple of 3 values
            else:
                error_message = "Decomposed SVG and Animation Suggestion not found in response."
                error_html = create_error_html(error_message)
                return "", "", error_html  # Return tuple of 3 values
        except Exception as e:
            error_message = f"Error during MLLM feedback prediction: {e}"
            error_html = create_error_html(error_message)
            return "", "", error_html  # Return tuple of 3 values

    def generate_animation(self, proposed_animation: str, svg_content: str) -> str:
        try:
            prompt = self.generate_animation_prompt.format(
                svg_content=svg_content, proposed_animation=proposed_animation
            )
            logger.info(f"Animation Generation Prompt:\n{prompt}")

            if self.client:
                response = self.client.messages.create(
                    model="claude-sonnet-4-20250514",
                    max_tokens=10000,
                    messages=[{"role": "user", "content": prompt}],
                )
                response_text = response.content[0].text
                logger.info(f"Model Response:\n{response_text}")
                return response_text

        except Exception as e:
            return f"<html><body><h3>Error generating animation: {e}</h3></body></html>"


generator = SVGAnimationGenerator()


def process_svg(svg_file):
    if svg_file is None:
        return "Please upload an SVG file"
    try:
        with open(svg_file, "r", encoding="utf-8") as f:
            svg_content = f.read()
        parsed_svg = generator.parse_svg(svg_content)
        return parsed_svg.get("svg_content", "")
    except FileNotFoundError:
        return "File not found. Please upload a valid SVG file."
    except ET.ParseError:
        return "Invalid SVG file format. Please upload a valid SVG file."
    except Exception as e:
        return f"Error processing file: {e}"


def predict_decompose_group(svg_file, svg_text, object_name):
    if not object_name.strip():
        error_msg = "Please enter a valid object name for the SVG"
        error_html = (error_msg)
        return "", error_msg, "", error_html

    if svg_file is not None:
        svg_content_inner = process_svg(svg_file)
    else:
        svg_content_inner = svg_text.strip()

    if not svg_content_inner:
        error_msg = "Please upload an SVG file or enter SVG markup"
        error_html = create_error_html(error_msg)
        return "", error_msg, "", error_html

    parsed_svg = generator.parse_svg(svg_content_inner)
    if "error" in parsed_svg:
        error_msg = parsed_svg["error"]
        error_html = create_error_html(error_msg)
        return "", error_msg, "", error_html

    decompose_result = generator.predict_decompose_group(parsed_svg, object_name)
    if "error" in decompose_result:
        error_msg = decompose_result["error"]
        error_html = create_error_html(error_msg)
        return "", error_msg, "", error_html

    decomposed_svg = decompose_result["svg_content"]
    animation_suggestions = decompose_result["animation_suggestions"]

    # Create viewer HTML
    decomposed_svg_viewer = f"""
    <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
        <div style='display: block; align-items: center; margin-bottom: 10px;'>
        </div>
        <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
            {decomposed_svg}
        </div>
    </div>
    """

    return (
        decomposed_svg,          # For svg_content_hidden
        decomposed_svg,             # For groups_summary (분석 결과 표시)
        animation_suggestions,   # For animation_suggestion
        decomposed_svg_viewer,   # For decomposed_svg_viewer
    )

def create_animation_preview(animation_desc: str, svg_content: str) -> str:
    """Create animation preview from description and SVG content."""
    if not svg_content.strip():
        return create_error_html("⚠️ Please process SVG first")

    if not animation_desc.strip():
        return create_error_html("⚠️ Please describe the animation you want")

    try:
        animation_html = generator.generate_animation(animation_desc, svg_content)
        if not animation_html:
            return create_error_html("❌ Failed to generate animation")

        # Extract HTML content from Claude's response
        html_match = re.search(
            r"<html_output>(.*?)</html_output>", animation_html, re.DOTALL
        )
        if not html_match:
            return create_error_html("❌ Invalid animation HTML format")

        # Get the actual HTML content
        html_content = html_match.group(1).strip()

        # Save the HTML content to the output directory
        output_dir = "output"
        os.makedirs(output_dir, exist_ok=True)
        html_path = os.path.join(output_dir, "animation_preview.html")
        with open(html_path, "w", encoding="utf-8") as f:
            f.write(html_content)
        print(f"Animation preview saved to: {html_path}")

        html_path = pathlib.Path("output/animation_preview.html").read_text(
            encoding="utf-8"
        )

        # Wrap in a container with preview styling
        return f"""
        <div style='padding: 20px; background: #fff; border: 1px solid #eee; border-radius: 8px; display: block;'>
            <div style='display: block; align-items: center; margin-bottom: 10px;'>
            </div>
            <div id='animation-container' style='min-height: 300px; display: block; justify-content: center; align-items: center; background: #fafafa; border-radius: 4px; padding: 20px;'>
            <iframe srcdoc="{html_path.replace('"', '&quot;')}"
                width="100%" height="600px" 
                style="border:none; overflow:hidden;"
                sandbox="allow-scripts allow-same-origin">
            </iframe>
            </div>
        </div>
        """
    except Exception as e:
        return create_error_html(f"❌ Error creating animation: {str(e)}")


def create_error_html(message: str) -> str:
    """Create formatted error message HTML."""
    return f"""
    <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
        <h3>{message}</h3>
    </div>
    """


# Define examples with proper path handling and categories
example_list = {
    "Animals": [
        [os.path.join(os.path.dirname(__file__), "examples/corgi.svg"), "corgi"],
        [os.path.join(os.path.dirname(__file__), "examples/duck.svg"), "duck"],
        [os.path.join(os.path.dirname(__file__), "examples/whale.svg"), "whale"],
    ],
    "Objects": [
        [os.path.join(os.path.dirname(__file__), "examples/rocket.svg"), "rocket"]
    ],
}


def load_example(example_choice):
    # Find the selected example in the categories
    for category, examples in example_list.items():
        for example in examples:
            if example[1] == example_choice:
                return example[0]  # Return the file path
    return None


# Flatten choices for dropdown
example_choices = [
    example[1] for category in example_list.values() for example in category
]

demo = gr.Blocks(title="SVG Animation Generator", theme=gr.themes.Soft())
with demo:
    gr.Markdown("# 🎨 SVG Decomposition & Animation Generator")
    gr.Markdown(
        "Intelligent SVG decomposition and animation generation powered by MLLM. This tool decomposes SVG structures and generates animations based on your descriptions."
    )
    with gr.Column():
        with gr.Row(scale=2):
            with gr.Column(scale=1):
                gr.Markdown("## 📤 Input SVG")
                with gr.Row(scale=2):
                    svg_file = gr.File(label="Upload SVG File", file_types=[".svg"])
                    svg_text = gr.Textbox(
                        label="Or Paste SVG Code",
                        lines=8.4,
                    )

                # Add example dropdown
                example_dropdown = gr.Dropdown(
                    choices=example_choices, label="Try an Example", value=None
                )

    # Add dropdown change event
    example_dropdown.change(
        fn=load_example, inputs=[example_dropdown], outputs=[svg_file]
    )

    with gr.Column(scale=1):
        with gr.Row(scale=1):
            with gr.Column(scale=1):
                gr.Markdown("## 🔍 SVG Analysis")
                object_name = gr.Textbox(
                    label="Name Your Object",
                    placeholder="Give a name to your SVG (e.g., 'dove', 'robot')",
                    value="corgi",
                )
                process_btn = gr.Button("🔄 Decompose Structure", variant="primary")
                groups_summary = gr.Textbox(
                    label="Decomposition Results",
                    placeholder="MLLM will analyze and decompose the SVG structure...",
                    lines=6,
                    interactive=False,
                )
            with gr.Column(scale=1):
                gr.Markdown("## 🎯 Animation Design")
                animation_suggestion = gr.Textbox(
                    label="AI Suggestions",
                    placeholder="MLLM will suggest animations based on the decomposed structure...",
                    lines=14.5,
                )
        with gr.Row(scale=1):
            with gr.Column(scale=1):
                gr.Markdown("## 💡 Decomposed Elements")
                groups_feedback = gr.Textbox(
                    label="Element Structure",
                    placeholder="If you have specific decomposition in mind, describe it here...",
                    lines=2,
                )
                groups_feedback_btn = gr.Button(
                    "💭 Apply Decomposition Feedback", variant="primary"
                )
            with gr.Column(scale=1):
                gr.Markdown("## ✨ Create Animation")
                describe_animation = gr.Textbox(
                    label="Animation Description",
                    placeholder="Describe your desired animation (e.g., 'gentle floating motion')",
                    lines=2,
                )
                animate_btn = gr.Button("🎬 Generate Animation", variant="primary")

        with gr.Row(scale=3):
            with gr.Column(scale=1):
                svg_content_hidden = gr.Textbox(visible=False)
                gr.Markdown("## 🖼️ Decomposed Structure")
                decomposed_svg_viewer = gr.HTML(
                    label="Decomposed SVG",
                    value="""
        <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
            <div id='decomposed-svg-container' style='min-height: 150px; display: flex; justify-content: center; align-items: center; border-radius: 4px; padding: 10px;'>
                <div style='color: #999; text-align: center;'>Decomposed SVG structure will appear here</div>
            </div>
        </div>
        """,
                )
            with gr.Column(scale=1):
                gr.Markdown("## 🎭 Animation Preview")
                animation_preview = gr.HTML(
                    label="Live Preview",
                    value="""
        <div style='padding: 40px; text-align: center; color: #666; border: 2px dashed #ddd; border-radius: 10px;'>
        <div id='animation-container' style='min-height: 150px; display: flex; justify-content: center; align-items: center; border-radius: 4px; padding: 10px;'>
                <div style='color: #999; text-align: center;'>Amination preview will appear here</div>
            </div>
        </div>
        """,
                )

    with gr.Column(): 
        with gr.Row(scale=1):
            with gr.Column(scale=1):
                gr.Markdown("## 📂 HTML Output")
                output_html = gr.Textbox(
                    label="Output HTML",
                    lines=10,
                    placeholder="Generated HTML will be saved here.",
                )

        process_btn.click(
        fn=predict_decompose_group,
        inputs=[svg_file, svg_text, object_name],
        outputs=[
            svg_content_hidden,    # Store decomposed SVG for later use
            groups_summary,        # Show analysis results
            animation_suggestion,  # Show animation suggestions
            decomposed_svg_viewer, # Show SVG preview
        ],
    )
    groups_feedback_btn.click(
        fn=generator.feedback_decompose_group,
        inputs=[
            svg_content_hidden,    # Pass the SVG content directly
            groups_feedback,       # Pass the feedback text
        ],
        outputs=[
            svg_content_hidden,    # Update hidden SVG content
            animation_suggestion,  # Update animation suggestions
            decomposed_svg_viewer, # Update SVG preview
        ],
    )
    animate_btn.click(
        fn=create_animation_preview,
        inputs=[
            describe_animation,
            svg_content_hidden,
        ],
        outputs=[
            animation_preview, 
            output_html
        ],
    )
if __name__ == "__main__":
    demo.launch(share=True)