"""
LLM 故事生成服务 — 基于火山引擎豆包大模型
"""
import json
import logging
from django.conf import settings

try:
    from volcenginesdkarkruntime import Ark
    ARK_AVAILABLE = True
except ImportError:
    ARK_AVAILABLE = False

logger = logging.getLogger(__name__)

STORY_SYSTEM_PROMPT = """# 角色

你是「卡皮巴拉故事工坊」的首席故事大师。你为 3-8 岁的小朋友创作原创童话故事。

# 任务

根据用户提供的**角色、场景、道具**素材，创作一个完整的儿童故事。

# 输出格式

你 **必须** 只返回如下 JSON，不要返回任何其他内容（不要 markdown 代码块，不要解释）：

{"title": "故事标题（6字以内）", "content": "故事正文"}

# 故事创作规范

1. **字数**：正文 400-600 字，不要太短也不要太长
2. **段落**：用 `\\n\\n` 分段，每段 2-4 句话
3. **语言**：简单易懂，适合给小朋友朗读；可以包含拟声词（"哗啦啦"、"咕噜噜"）和语气词（"哇！"、"嘿嘿"）
4. **结构**：开头引入角色和场景 → 中间遇到挑战或趣事 → 结尾温馨圆满
5. **情感**：温暖、有趣、充满想象力，带一点小幽默
6. **教育**：自然融入一个小道理（勇气、友谊、分享等），不要说教
7. **创意**：即使收到相同的素材组合，每次也要创作全新的、不同的故事情节
8. **角色融合**：所有用户选择的角色、场景、道具都必须在故事中出现并发挥作用
9. **标题**：简短有趣，6 个字以内，能引起小朋友的好奇心"""


def build_user_prompt(characters, scenes, props):
    """构建用户提示词"""
    parts = []
    if characters:
        parts.append(f"角色：{', '.join(characters)}")
    if scenes:
        parts.append(f"场景：{', '.join(scenes)}")
    if props:
        parts.append(f"道具：{', '.join(props)}")
    return '请根据以下元素创作一个儿童故事：\n' + '\n'.join(parts)


def generate_story_stream(characters, scenes, props):
    """
    流式生成故事，yield SSE 事件字符串。
    使用火山引擎豆包大模型（Ark SDK）。

    Yields:
        str: SSE 格式的事件数据行
    """
    config = settings.LLM_CONFIG

    if not config.get('API_KEY'):
        yield _sse_event('error', {'message': 'Volcengine API Key 未配置'})
        return

    if not ARK_AVAILABLE:
        yield _sse_event('error', {'message': 'volcengine SDK 未安装，请运行 pip install "volcengine-python-sdk[ark]"'})
        return

    yield _sse_event('stage', {
        'stage': 'connecting',
        'progress': 0,
        'message': '正在收集灵感碎片...',
    })

    client = Ark(api_key=config['API_KEY'])

    user_prompt = build_user_prompt(characters, scenes, props)

    try:
        yield _sse_event('stage', {
            'stage': 'generating',
            'progress': 10,
            'message': '故事正在诞生...',
        })

        stream = client.chat.completions.create(
            model=config['MODEL_NAME'],
            messages=[
                {'role': 'system', 'content': STORY_SYSTEM_PROMPT},
                {'role': 'user', 'content': user_prompt},
            ],
            max_tokens=2048,
            stream=True,
        )

        full_content = ''
        chunk_count = 0

        for chunk in stream:
            delta = chunk.choices[0].delta if chunk.choices else None
            if delta and delta.content:
                full_content += delta.content
                chunk_count += 1

                if chunk_count % 5 == 0:
                    progress = min(10 + int(chunk_count * 0.5), 80)
                    yield _sse_event('stage', {
                        'stage': 'generating',
                        'progress': progress,
                        'message': '故事正在诞生...',
                    })

        yield _sse_event('stage', {
            'stage': 'parsing',
            'progress': 85,
            'message': '正在编制最后的魔法...',
        })

        result = _parse_story_json(full_content)

        yield _sse_event('done', {
            'stage': 'done',
            'progress': 100,
            'message': '大功告成！',
            'title': result['title'],
            'content': result['content'],
        })

    except Exception as e:
        logger.error(f'LLM story generation failed: {e}')
        yield _sse_event('error', {'message': f'故事生成失败: {str(e)}'})


def _parse_story_json(text):
    """从 LLM 输出中解析故事 JSON"""
    text = text.strip()
    if text.startswith('```'):
        text = text.split('\n', 1)[1] if '\n' in text else text[3:]
    if text.endswith('```'):
        text = text[:-3]
    text = text.strip()

    try:
        data = json.loads(text)
        return {
            'title': data.get('title', '未命名故事'),
            'content': data.get('content', text),
        }
    except json.JSONDecodeError:
        return {
            'title': '新故事',
            'content': text,
        }


def _sse_event(event, data):
    """格式化 SSE 事件"""
    return f"event: {event}\ndata: {json.dumps(data, ensure_ascii=False)}\n\n"