from typing import List, Optional, Dict
import json
from .language_config import (
    LanguageType,
    DEFAULT_CATEGORIES,
    CATEGORY_DESCRIPTIONS,
)


def get_default_categories(lang: LanguageType = "zh") -> List[str]:
    return DEFAULT_CATEGORIES.get(lang, DEFAULT_CATEGORIES["zh"])


def get_category_descriptions(lang: LanguageType = "zh") -> Dict[str, str]:
    return CATEGORY_DESCRIPTIONS.get(lang, CATEGORY_DESCRIPTIONS["zh"])


def get_attribute_prompt(
    query: str,
    categories: Optional[List[str]] = None,
    lang: LanguageType = "zh"
) -> str:
    """Generate prompt with causal chain structure."""
    if lang == "en":
        prompt = f"""Analyze the attributes of "{query}" in a causal chain format: Materials→Functions→Usages→User Groups.

List 3-5 types of materials, each extending into a complete causal chain.

JSON format:
{{"name": "{query}", "children": [{{"name": "Material Name", "category": "Materials", "children": [{{"name": "Function Name", "category": "Functions", "children": [{{"name": "Usage Name", "category": "Usages", "children": [{{"name": "User Group Name", "category": "User Groups"}}]}}]}}]}}]}}

Return JSON only."""
    else:
        prompt = f"""分析「{query}」的屬性，以因果鏈方式呈現：材料→功能→用途→使用族群。

請列出 3-5 種材料，每種材料延伸出完整因果鏈。

JSON 格式：
{{"name": "{query}", "children": [{{"name": "材料名", "category": "材料", "children": [{{"name": "功能名", "category": "功能", "children": [{{"name": "用途名", "category": "用途", "children": [{{"name": "族群名", "category": "使用族群"}}]}}]}}]}}]}}

只回傳 JSON。"""

    return prompt


def get_step1_attributes_prompt(query: str, lang: LanguageType = "zh") -> str:
    """Step 1: Generate attribute list for each category (parallel structure)"""
    if lang == "en":
        return f"""/no_think
Analyze "{query}" and list attributes for the following four categories. List 3-5 common attributes for each category.

Return JSON only, in the following format:
{{"materials": ["material1", "material2", "material3"], "functions": ["function1", "function2", "function3"], "usages": ["usage1", "usage2", "usage3"], "users": ["user group1", "user group2", "user group3"]}}

Object: {query}"""
    else:
        return f"""/no_think
分析「{query}」，列出以下四個類別的屬性。每個類別列出 3-5 個常見屬性。

只回傳 JSON，格式如下：
{{"materials": ["材料1", "材料2", "材料3"], "functions": ["功能1", "功能2", "功能3"], "usages": ["用途1", "用途2", "用途3"], "users": ["族群1", "族群2", "族群3"]}}

物件：{query}"""


def get_step2_causal_chain_prompt(
    query: str,
    materials: List[str],
    functions: List[str],
    usages: List[str],
    users: List[str],
    existing_chains: List[dict],
    chain_index: int,
    lang: LanguageType = "zh"
) -> str:
    """Step 2: Generate a single causal chain"""
    existing_chains_text = ""

    if lang == "en":
        if existing_chains:
            chains_list = [
                f"- {c['material']} → {c['function']} → {c['usage']} → {c['user']}"
                for c in existing_chains
            ]
            existing_chains_text = f"""
[Already generated causal chains, do not repeat]
{chr(10).join(chains_list)}
"""
        return f"""/no_think
Generate causal chain #{chain_index} for "{query}".

[Available Materials] {', '.join(materials)}
[Available Functions] {', '.join(functions)}
[Available Usages] {', '.join(usages)}
[Available User Groups] {', '.join(users)}
{existing_chains_text}
[Rules]
1. Select one attribute from each category to form a logical causal chain
2. The causal relationship must be logical (materials determine functions, functions determine usages, usages determine user groups)
3. Do not repeat existing causal chains

Return JSON only:
{{"material": "selected material", "function": "selected function", "usage": "selected usage", "user": "selected user group"}}"""
    else:
        if existing_chains:
            chains_list = [
                f"- {c['material']} → {c['function']} → {c['usage']} → {c['user']}"
                for c in existing_chains
            ]
            existing_chains_text = f"""
【已生成的因果鏈，請勿重複】
{chr(10).join(chains_list)}
"""
        return f"""/no_think
為「{query}」生成第 {chain_index} 條因果鏈。

【可選材料】{', '.join(materials)}
【可選功能】{', '.join(functions)}
【可選用途】{', '.join(usages)}
【可選族群】{', '.join(users)}
{existing_chains_text}
【規則】
1. 從每個類別選擇一個屬性，組成合理的因果鏈
2. 因果關係必須合邏輯（材料決定功能，功能決定用途，用途決定族群）
3. 不要與已生成的因果鏈重複

只回傳 JSON：
{{"material": "選擇的材料", "function": "選擇的功能", "usage": "選擇的用途", "user": "選擇的族群"}}"""


def get_flat_attribute_prompt(
    query: str,
    categories: Optional[List[str]] = None,
    lang: LanguageType = "zh"
) -> str:
    """Generate prompt with flat/parallel categories (original design)."""
    cats = categories if categories else get_default_categories(lang)
    cat_descs = get_category_descriptions(lang)

    # Build category list
    category_lines = []
    for cat in cats:
        desc = cat_descs.get(cat, f"Related attributes of {cat}" if lang == "en" else f"{cat}的相關屬性")
        category_lines.append(f"- {cat}: {desc}")

    categories_text = "\n".join(category_lines)

    if lang == "en":
        prompt = f"""/no_think
You are an object attribute analysis expert. Please break down the user's input object into the following attribute categories.

[Required Categories]
{categories_text}

[Important] The return format must be valid JSON, and each node must have a "name" field:

```json
{{
  "name": "Object Name",
  "children": [
    {{
      "name": "Category Name",
      "children": [
        {{"name": "Attribute 1"}},
        {{"name": "Attribute 2"}}
      ]
    }}
  ]
}}
```

Return JSON only, no other text.

User input: {query}"""
    else:
        prompt = f"""/no_think
你是一個物件屬性分析專家。請將用戶輸入的物件拆解成以下屬性類別。

【必須包含的類別】
{categories_text}

【重要】回傳格式必須是合法的 JSON，每個節點都必須有 "name" 欄位：

```json
{{
  "name": "物件名稱",
  "children": [
    {{
      "name": "類別名稱",
      "children": [
        {{"name": "屬性1"}},
        {{"name": "屬性2"}}
      ]
    }}
  ]
}}
```

只回傳 JSON，不要有任何其他文字。

用戶輸入：{query}"""

    return prompt


# ===== Dynamic category system prompts =====

def get_step0_category_analysis_prompt(
    query: str,
    suggested_count: int = 3,
    exclude_categories: List[str] | None = None,
    lang: LanguageType = "zh"
) -> str:
    """Step 0: LLM analyzes and suggests categories"""

    if lang == "en":
        exclude_text = ""
        if exclude_categories:
            exclude_text = f"\n[Forbidden Categories] {', '.join(exclude_categories)} (These are already fixed categories, do not suggest duplicates)\n"

        return f"""/no_think
Analyze "{query}" and suggest {suggested_count} most suitable attribute categories to describe it.

[Common Category References] Characteristics, Shape, Color, Size, Brand, Price Range, Weight, Style, Occasion, Season, Technical Specifications
{exclude_text}
[Important]
1. Choose categories that best describe the essence of this object
2. Categories should have logical relationships
3. Do not choose overly abstract or duplicate categories
4. Must suggest creative categories different from the reference list

Return JSON only:
{{
  "categories": [
    {{"name": "Category1", "description": "Description1", "order": 0}},
    {{"name": "Category2", "description": "Description2", "order": 1}}
  ]
}}

Object: {query}"""
    else:
        exclude_text = ""
        if exclude_categories:
            exclude_text = f"\n【禁止使用的類別】{', '.join(exclude_categories)}（這些已經是固定類別，不要重複建議）\n"

        return f"""/no_think
分析「{query}」，建議 {suggested_count} 個最適合的屬性類別來描述它。

【常見類別參考】特性、形狀、顏色、尺寸、品牌、價格區間、重量、風格、場合、季節、技術規格
{exclude_text}
【重要】
1. 選擇最能描述此物件本質的類別
2. 類別之間應該有邏輯關係
3. 不要選擇過於抽象或重複的類別
4. 必須建議與參考列表不同的、有創意的類別

只回傳 JSON：
{{
  "categories": [
    {{"name": "類別1", "description": "說明1", "order": 0}},
    {{"name": "類別2", "description": "說明2", "order": 1}}
  ]
}}

物件：{query}"""


def get_step1_dynamic_attributes_prompt(
    query: str,
    categories: List,  # List[CategoryDefinition]
    lang: LanguageType = "zh"
) -> str:
    """Dynamic Step 1 - Generate attributes based on category list"""
    # Sort by order and build description
    sorted_cats = sorted(categories, key=lambda x: x.order if hasattr(x, 'order') else x.get('order', 0))

    category_desc = "\n".join([
        f"- {cat.name if hasattr(cat, 'name') else cat['name']}: {cat.description if hasattr(cat, 'description') else cat.get('description', 'Related attributes' if lang == 'en' else '相關屬性')}"
        for cat in sorted_cats
    ])

    category_keys = [cat.name if hasattr(cat, 'name') else cat['name'] for cat in sorted_cats]

    if lang == "en":
        json_template = {cat: ["attribute1", "attribute2", "attribute3"] for cat in category_keys}
        return f"""/no_think
Analyze "{query}" and list attributes for the following categories. List 3-5 common attributes for each category.

[Category List]
{category_desc}

Return JSON only:
{json.dumps(json_template, ensure_ascii=False, indent=2)}

Object: {query}"""
    else:
        json_template = {cat: ["屬性1", "屬性2", "屬性3"] for cat in category_keys}
        return f"""/no_think
分析「{query}」，列出以下類別的屬性。每個類別列出 3-5 個常見屬性。

【類別列表】
{category_desc}

只回傳 JSON：
{json.dumps(json_template, ensure_ascii=False, indent=2)}

物件：{query}"""


def get_step2_dynamic_causal_chain_prompt(
    query: str,
    categories: List,  # List[CategoryDefinition]
    attributes_by_category: Dict[str, List[str]],
    existing_chains: List[Dict[str, str]],
    chain_index: int,
    lang: LanguageType = "zh"
) -> str:
    """Dynamic Step 2 - Generate causal chains for dynamic categories"""
    sorted_cats = sorted(categories, key=lambda x: x.order if hasattr(x, 'order') else x.get('order', 0))

    # Build available attributes
    available_attrs = "\n".join([
        f"[{cat.name if hasattr(cat, 'name') else cat['name']}] {', '.join(attributes_by_category.get(cat.name if hasattr(cat, 'name') else cat['name'], []))}"
        for cat in sorted_cats
    ])

    if lang == "en":
        # Already generated causal chains
        existing_text = ""
        if existing_chains:
            chains_list = [
                " → ".join([chain.get(cat.name if hasattr(cat, 'name') else cat['name'], '?') for cat in sorted_cats])
                for chain in existing_chains
            ]
            existing_text = "\n[Already generated, do not repeat]\n" + "\n".join([f"- {c}" for c in chains_list])

        # JSON template
        json_template = {cat.name if hasattr(cat, 'name') else cat['name']: f"selected {cat.name if hasattr(cat, 'name') else cat['name']}" for cat in sorted_cats}

        return f"""/no_think
Generate causal chain #{chain_index} for "{query}".

[Available Attributes]
{available_attrs}
{existing_text}

[Rules]
1. Select one attribute from each category
2. Causal relationships must be logical
3. Do not repeat

Return JSON only:
{json.dumps(json_template, ensure_ascii=False, indent=2)}"""
    else:
        # 已生成的因果鏈
        existing_text = ""
        if existing_chains:
            chains_list = [
                " → ".join([chain.get(cat.name if hasattr(cat, 'name') else cat['name'], '?') for cat in sorted_cats])
                for chain in existing_chains
            ]
            existing_text = "\n【已生成，請勿重複】\n" + "\n".join([f"- {c}" for c in chains_list])

        # JSON 模板
        json_template = {cat.name if hasattr(cat, 'name') else cat['name']: f"選擇的{cat.name if hasattr(cat, 'name') else cat['name']}" for cat in sorted_cats}

        return f"""/no_think
為「{query}」生成第 {chain_index} 條因果鏈。

【可選屬性】
{available_attrs}
{existing_text}

【規則】
1. 從每個類別選擇一個屬性
2. 因果關係必須合理
3. 不要重複

只回傳 JSON：
{json.dumps(json_template, ensure_ascii=False, indent=2)}"""


# ===== DAG relationship prompt =====

def get_step2_dag_relationships_prompt(
    query: str,
    categories: List,  # List[CategoryDefinition]
    attributes_by_category: Dict[str, List[str]],
    lang: LanguageType = "zh"
) -> str:
    """Generate natural relationships between adjacent categories"""
    sorted_cats = sorted(categories, key=lambda x: x.order if hasattr(x, 'order') else x.get('order', 0))

    # Build attribute listing
    attr_listing = "\n".join([
        f"[{cat.name if hasattr(cat, 'name') else cat['name']}] {', '.join(attributes_by_category.get(cat.name if hasattr(cat, 'name') else cat['name'], []))}"
        for cat in sorted_cats
    ])

    # Build direction hints
    direction_hints = " → ".join([cat.name if hasattr(cat, 'name') else cat['name'] for cat in sorted_cats])

    if lang == "en":
        return f"""/no_think
Analyze the attribute relationships of "{query}".

{attr_listing}

[Relationship Direction] {direction_hints}

[Rules]
1. Only establish relationships between adjacent categories (e.g., Materials→Functions, Functions→Usages)
2. Only output pairs that have true causal or associative relationships
3. An attribute can connect to multiple downstream attributes, or none at all
4. Not every attribute needs to have connections
5. Relationships should be reasonable and meaningful

Return JSON:
{{
  "relationships": [
    {{"source_category": "CategoryA", "source": "attribute name", "target_category": "CategoryB", "target": "attribute name"}},
    ...
  ]
}}

Return JSON only."""
    else:
        return f"""/no_think
分析「{query}」的屬性關係。

{attr_listing}

【關係方向】{direction_hints}

【規則】
1. 只建立相鄰類別之間的關係（例如：材料→功能，功能→用途）
2. 只輸出真正有因果或關聯關係的配對
3. 一個屬性可連接多個下游屬性，也可以不連接任何屬性
4. 不需要每個屬性都有連接
5. 關係應該合理且有意義

回傳 JSON：
{{
  "relationships": [
    {{"source_category": "類別A", "source": "屬性名", "target_category": "類別B", "target": "屬性名"}},
    ...
  ]
}}

只回傳 JSON。"""
