import json
import requests
import re
from pathlib import Path
from typing import Dict, List, Optional

from prompts import get_prompt_download_assets, get_prompt_place_assets


class SmartSVGDownloader:
    def __init__(self, assets_dir: str, api_function=None, iconfinder_api_key: str = None):
        self.assets_dir = Path(assets_dir)
        self.assets_dir.mkdir(exist_ok=True)
        self.api_function = api_function
        self.iconfinder_api_key = iconfinder_api_key

    def process_storyboard(self, storyboard: Dict) -> Dict:
        storyboard_data = json.loads(json.dumps(storyboard))
        sections = storyboard_data.get("sections", [])
        selected_sections = []
        if sections:
            selected_sections.append(sections[0])
            if len(sections) > 1:
                selected_sections.append(sections[-1])
        temp_storyboard = {"sections": selected_sections}
        # print(temp_storyboard)

        elements = self._analyze_assets_needed(temp_storyboard)

        # First, check the local cache. Only download what is missing
        downloaded_assets = {}
        for el in elements:
            cached = self._check_cache(el)
            if cached:
                downloaded_assets[el] = cached
            else:
                filepath = self._download_element(el)
                if filepath:
                    downloaded_assets[el] = filepath
                    print(f"✓ 下载: {el} -> {filepath}")

        prompt = self._build_enhancement_prompt(storyboard, downloaded_assets)
        api_response = self.api_function(prompt, max_tokens=2000)[0]

        enhanced_storyboard = self._parse_api_response(api_response, storyboard_data)
        return enhanced_storyboard

    def _build_enhancement_prompt(self, storyboard: Dict, downloaded_assets: Dict) -> str:
        asset_mapping = ""
        if downloaded_assets:
            asset_mapping = "Available Assets:\n"
            for element, filepath in downloaded_assets.items():
                asset_mapping += f"- {element}: [Asset: {filepath}]\n"
            asset_mapping += "\n"
        sections = storyboard.get("sections", [])
        animations_data = []
        if sections:
            first = sections[0]
            animations_data.append(
                {"section_index": 0, "section_id": first.get("id", ""), "animations": first.get("animations", [])}
            )
            if len(sections) > 1:
                last = sections[-1]
                animations_data.append(
                    {
                        "section_index": len(sections) - 1,
                        "section_id": last.get("id", ""),
                        "animations": last.get("animations", []),
                    }
                )
        animations_structure = json.dumps(animations_data, indent=2, ensure_ascii=False)
        return get_prompt_place_assets(asset_mapping, animations_structure)

    def _extract_json_from_markdown(self, text: str) -> str:
        pattern = r"```(?:json)?\s*([\{\[].*?[\}\]])\s*```"
        m = re.search(pattern, text, re.DOTALL)
        return m.group(1) if m else text

    def _parse_api_response(self, response: str, original_storyboard: Dict) -> Dict:
        """Parse API response and update storyboard"""
        try:
            try:
                content = response.candidates[0].content.parts[0].text
            except Exception:
                try:
                    content = response.choices[0].message.content
                except Exception:
                    content = str(response)

            enhanced_animations = json.loads(self._extract_json_from_markdown(content))

            # Create a copy of the storyboard for enhancement
            enhanced_storyboard = json.loads(json.dumps(original_storyboard))

            if isinstance(enhanced_animations, list):
                for anim_data in enhanced_animations:
                    section_index = anim_data.get("section_index")
                    enhanced_anims = anim_data.get("animations", [])

                    if isinstance(section_index, int) and 0 <= section_index < len(enhanced_storyboard.get("sections", [])):
                        enhanced_storyboard["sections"][section_index]["animations"] = enhanced_anims

            return enhanced_storyboard

        except json.JSONDecodeError as e:
            print(f"API response parsing failed: {e}")
            return original_storyboard
        except Exception as e:
            print(f"Error occurred while processing API response: {e}")
            return original_storyboard

    def _analyze_assets_needed(self, storyboard_data) -> List[str]:
        if not storyboard_data:
            return []

        prompt = get_prompt_download_assets(storyboard_data=storyboard_data)
        try:
            response = self.api_function(prompt, max_tokens=100)[0]
            try:
                content = response.candidates[0].content.parts[0].text
            except:
                content = response.choices[0].message.content
            elements = [line.strip().lower() for line in content.strip().split("\n") if line.strip()]
            return list(dict.fromkeys(elements))[:4]
        except:
            return []

    def _check_cache(self, element: str) -> Optional[str]:
        for suffix in [".png", ".svg"]:
            filepath = self.assets_dir / f"{element}{suffix}"
            if filepath.exists():
                return str(filepath.absolute())
        return None

    def _download_element(self, element: str) -> Optional[str]:
        return self._download_iconfinder(element) or self._download_iconify(element)

    def _download_iconfinder(self, element: str) -> Optional[str]:
        try:
            url = f"https://api.iconfinder.com/v4/icons/search?query={element}&count=1&premium=0"
            headers = {"Authorization": f"Bearer {self.iconfinder_api_key}"}
            resp = requests.get(url, headers=headers, timeout=10)
            if resp.status_code != 200:
                return None
            data = resp.json()
            if not data.get("icons"):
                return None
            raster_sizes = data["icons"][0].get("raster_sizes", [])
            size_url = None
            for size in [256, 128, 512]:
                for s in raster_sizes:
                    if s["size"] == size:
                        size_url = s["formats"][0]["preview_url"]
                        break
                if size_url:
                    break
            if not size_url and raster_sizes:
                size_url = raster_sizes[-1]["formats"][0]["preview_url"]
            if size_url:
                img_resp = requests.get(size_url, timeout=10)
                if img_resp.status_code == 200:
                    filepath = self.assets_dir / f"{element}.png"
                    filepath.write_bytes(img_resp.content)
                    return str(filepath.absolute())
        except:
            return None

    def _download_iconify(self, element: str) -> Optional[str]:
        try:
            search_url = f"https://api.iconify.design/search?query={element}&limit=1"
            r = requests.get(search_url, timeout=8)
            if r.status_code == 200 and r.json().get("icons"):
                icon_id = r.json()["icons"][0]
                collection, name = icon_id.split(":", 1)
                svg_url = f"https://api.iconify.design/{collection}/{name}.svg"
                svg_resp = requests.get(svg_url, timeout=8)
                if svg_resp.status_code == 200:
                    filepath = self.assets_dir / f"{element}.svg"
                    filepath.write_text(svg_resp.text, encoding="utf-8")
                    return str(filepath.absolute())
        except:
            return None

    def _enhance_animations(self, animations: List[str], assets: Dict[str, str]) -> List[str]:
        new_animations = []
        for anim in animations:
            for el, path in assets.items():
                if el in anim.lower() and path not in anim:
                    anim += f" [Asset: {path}]"
            new_animations.append(anim)
        return new_animations


def process_storyboard_with_assets(
    storyboard: Dict, api_function, assets_dir: str = "./assets/icon", iconfinder_api_key: str = None
) -> Dict:
    downloader = SmartSVGDownloader(assets_dir, api_function, iconfinder_api_key)
    return downloader.process_storyboard(storyboard)


if __name__ == "__main__":
    from gpt_request import request_gpt41_token

    sb = {
        "sections": [
            {
                "lecture_lines": ["A robot will guide the lesson", "The computer will process the data"],
                "animations": ["Show robot", "Display computer screen"],
            },
            {
                "lecture_lines": ["We will draw circles"],
                "animations": ["Draw blue circles"],
            },
        ]
    }

    downloader = SmartSVGDownloader("./assets/icon", request_gpt41_token, "Your API token")
    result = downloader.process_storyboard(sb)
    print(json.dumps(result, indent=2, ensure_ascii=False))
