import json


# Complete transformation examples used for few-shot learning
ALL_TRANSFORMATION_EXAMPLES = {
    "socrates is human": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "predicate": {"relation": "is", "complement": "human"},
        "certainty": "certain",
    },
    "socrates is mortal": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "predicate": {"relation": "is", "complement": "mortal"},
        "certainty": "certain",
    },
    "socrates sells iphone 12": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "predicate": {"relation": "sells", "complement": "iphone 12"},
        "certainty": "certain",
    },
    "socrates sells more than 20 iphones": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "quantity": "more than 20",
        "predicate": {"relation": "sells", "complement": "iphone"},
        "certainty": "certain",
    },
    "socrates sells iphone 15 and iphone 16": {
        "sentences": [
            {
                "intent_type": "making_statement",
                "subject": "socrates",
                "predicate": {"relation": "sells", "complement": "iphone 15"},
                "certainty": "certain",
            },
            {
                "intent_type": "making_statement",
                "subject": "socrates",
                "predicate": {"relation": "sells", "complement": "iphone 16"},
                "certainty": "certain",
            },
        ]
    },
    "iphone 12 is device": {
        "intent_type": "making_statement",
        "subject": "iphone 12",
        "predicate": {"relation": "is", "complement": "device"},
        "certainty": "certain",
    },
    "iphone 12 is form of device": {
        "intent_type": "making_statement",
        "subject": "iphone 12",
        "predicate": {"relation": "is form of", "complement": "device"},
        "certainty": "certain",
    },
    "iphone 12 has 13 inch screen size": {
        "sentences": [
            {
                "intent_type": "making_notion_attribute_statement",
                "subject": {
                    "caption": "iphone 12",
                    "attributes": [{"tag": "size", "value": 13, "unit": "inch"}],
                },
                "certainty": "certain",
            },
            {
                "intent_type": "making_statement",
                "subject": "iphone 12",
                "predicate": {"relation": "has", "complement": "screen"},
                "certainty": "certain",
            },
        ]
    },
    "iphone 12 costs 800 usd": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "iphone 12",
            "attributes": [{"tag": "price", "value": 800, "unit": "usd"}],
        },
        "certainty": "certain",
    },
    "iphone 12 has release year 2020": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "iphone 12",
            "attributes": [{"tag": "release year", "value": 2020}],
        },
        "certainty": "certain",
    },
    "iphone 12 is by socrates 1975": {
        "intent_type": "making_statement",
        "subject": "iphone 12",
        "predicate": {"relation": "is by", "complement": "socrates 1975"},
        "certainty": "certain",
    },
    "iphone 12 was directed by socrates": {
        "intent_type": "making_statement",
        "subject": "iphone 12",
        "predicate": {"relation": "directed by", "complement": "socrates"},
        "certainty": "certain",
    },
    "iphone 12 was produced by socrates": {
        "intent_type": "making_statement",
        "subject": "iphone 12",
        "predicate": {"relation": "produced by", "complement": "socrates"},
        "certainty": "certain",
    },
    "socrates birth year is 1970": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "socrates",
            "attributes": [{"tag": "birth year", "value": 1970}],
        },
        "certainty": "certain",
    },
    "socrates death year is 2001": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "socrates",
            "attributes": [{"tag": "death year", "value": 2001}],
        },
        "certainty": "certain",
    },
    "socrates was born on january 28, 1864": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "socrates",
            "attributes": [{"tag": "birth date", "value": "January 28, 1864"}],
        },
        "certainty": "certain",
    },
    "iphone 12 color is red": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "iphone 12",
            "attributes": [{"tag": "color", "value": "red"}],
        },
        "certainty": "certain",
    },
    "iphone 12 depth is 7 mm": {
        "intent_type": "making_notion_attribute_statement",
        "subject": {
            "caption": "iphone 12",
            "attributes": [{"tag": "depth", "value": 7, "unit": "mm"}],
        },
        "certainty": "certain",
    },
    "socrates sold iphone 12 in 2020": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "predicate": {"relation": "sold", "complement": "iphone 12"},
        "time": "2020",
        "certainty": "certain",
    },
    "socrates will sell iphone 12 in 2025": {
        "intent_type": "making_statement",
        "subject": "socrates",
        "predicate": {"relation": "will sell", "complement": "iphone 12"},
        "time": "2025",
        "certainty": "certain",
    },
    "philosophy magazine (1844–1846) was an american literary periodical": {
        "sentences": [
            {
                "intent_type": "making_notion_attribute_statement",
                "subject": {
                    "caption": "philosophy magazine",
                    "attributes": [
                        {"tag": "start year", "value": 1844},
                        {"tag": "end year", "value": 1846},
                    ],
                },
                "certainty": "certain",
            },
            {
                "intent_type": "making_statement",
                "subject": "philosophy magazine",
                "predicate": {"relation": "is", "complement": "an american literary periodical"},
                "certainty": "certain",
            },
        ]
    },
    "if socrates is human, then socrates is mortal": {
        "intent_type": "making_hypothetical_statement",
        "causes": [
            {
                "subject": "socrates",
                "certainty": "certain",
                "predicate": {"relation": "is", "complement": "human"},
            }
        ],
        "effects": [
            {
                "subject": "socrates",
                "certainty": "certain",
                "predicate": {"relation": "is", "complement": "mortal"},
            }
        ],
    },
}


def _apply_notion_attribute_schema(obj):
    if isinstance(obj, list):
        return [_apply_notion_attribute_schema(item) for item in obj]
    if not isinstance(obj, dict):
        return obj

    if "sentences" in obj and isinstance(obj["sentences"], list):
        obj = dict(obj)
        obj["sentences"] = [_apply_notion_attribute_schema(item) for item in obj["sentences"]]
        return obj

    obj = {k: _apply_notion_attribute_schema(v) for k, v in obj.items()}

    attributes = obj.get("attributes")
    subject = obj.get("subject")
    if attributes and isinstance(subject, str):
        obj = dict(obj)
        obj.pop("attributes", None)
        obj["subject"] = {"caption": subject, "attributes": attributes}

    return obj


def get_all_examples():
    """Get all transformation examples as a dictionary of JSON strings."""
    return {k: json.dumps(_apply_notion_attribute_schema(v)) for k, v in ALL_TRANSFORMATION_EXAMPLES.items()}


def build_few_shot_examples(max_examples: int | None = None) -> str:
    """Build formatted few-shot examples for LLM prompts."""
    examples = []
    items = list(ALL_TRANSFORMATION_EXAMPLES.items())
    if max_examples is not None:
        items = items[:max_examples]
    for input_text, output_json in items:
        normalized = _apply_notion_attribute_schema(output_json)
        examples.append(f"Input: {input_text}\nOutput: {normalized}\n")
    return "\n".join(examples)


def get_enhanced_system_prompt(base_prompt: str, num_examples: int | None = None) -> str:
    examples = build_few_shot_examples(num_examples)
    if "__FEW_SHOT_EXAMPLES__" in base_prompt:
        return base_prompt.replace("__FEW_SHOT_EXAMPLES__", examples)
    return base_prompt + "\n\n" + examples
