from tools.generate import generate_image
from tools.edit import edit_image
from tools.analyze import analyze_image
from tools.upload import save_uploads


class ImageAgent:

    def __init__(self, router, memory, planner):
        self.router = router
        self.memory = memory
        self.planner = planner

    # --------------------------------------------------
    # Upload Files
    # --------------------------------------------------

    def upload(self, files):

        uploaded = save_uploads(files)

        self.memory.uploaded_files = uploaded

        self.memory.image_context = None
        self.memory.analysis = None

        for file in uploaded:

            if file.get("type", "").startswith("image/"):

                self.memory.last_used_image = file["path"]

                break

        return uploaded

    # --------------------------------------------------
    # Analyze Image
    # --------------------------------------------------

    def auto_analyze(self, image_path=None):

        image = image_path or self.memory.last_used_image

        if not image:

            return {
                "description": "No uploaded image found."
            }

        analysis = analyze_image(image)

        self.memory.analysis = analysis
        self.memory.image_context = analysis

        return analysis

    # --------------------------------------------------
    # Convert object list safely
    # --------------------------------------------------

    def format_objects(self, objects):

        if not objects:
            return ""

        names = []

        for obj in objects:

            if isinstance(obj, str):

                names.append(obj)

            elif isinstance(obj, dict):

                if "name" in obj:

                    names.append(str(obj["name"]))

                elif "label" in obj:

                    names.append(str(obj["label"]))

                else:

                    names.append(str(obj))

            else:

                names.append(str(obj))

        return ", ".join(names)

    # --------------------------------------------------
    # Build Prompt
    # --------------------------------------------------

    def build_context(self, analysis):

        if not analysis:
            return ""

        objects = self.format_objects(
            analysis.get("objects", [])
        )

        colors = analysis.get("colors", [])

        if isinstance(colors, list):
            colors = ", ".join(colors)

        suggestions = analysis.get("suggestions", [])

        if isinstance(suggestions, list):
            suggestions = ", ".join(suggestions)

        context = f"""
Image Description:
{analysis.get("description","")}

Objects:
{objects}

People:
{analysis.get("people","")}

Style:
{analysis.get("style","")}

Lighting:
{analysis.get("lighting","")}

Colors:
{colors}

Background:
{analysis.get("background","")}

Camera Angle:
{analysis.get("camera_angle","")}

Quality:
{analysis.get("quality","")}

Editing Prompt:
{analysis.get("editing_prompt","")}

Suggestions:
{suggestions}
"""

        return context

    # --------------------------------------------------
    # Main Agent
    # --------------------------------------------------

    def run(self, user_input):

        task = self.router.route(user_input)

        steps = self.planner.plan(task)

        image = self.memory.last_used_image

        if image and self.memory.image_context is None:

            self.auto_analyze(image)

        analysis = self.memory.image_context

        context = self.build_context(analysis)

        prompt = f"""
{context}

User Request:
{user_input}
"""

        result = None

        for step in steps:

            if step == "generate":

                result = generate_image(prompt)

            elif step == "edit":

                result = edit_image(
                    image_path=image,
                    prompt=prompt
                )

                # If edit_image returns a file path
                if isinstance(result, str):

                    self.memory.last_used_image = result

                # If edit_image returns a PIL.Image,
                # app.py will display it directly.

            elif step == "analyze":

                result = self.auto_analyze(image)

            elif step == "chat":

                result = analysis

        return result