#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""小说2 -> 画布导入文件（分阶段断点续跑）
用法: python3 -u novel2_import.py [小说文件] [target_chars]
产出: aicomic-import.json  断点: novel2_import_ckpt.json
"""
import json, os, sys, urllib.request

BASE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, BASE)
import aicomic_import as ai

CKPT = os.path.join(BASE, 'novel2_import_ckpt.json')
OUT = os.path.join(BASE, 'aicomic-import.json')
TEXT = sys.argv[1] if len(sys.argv) > 1 else os.path.join(BASE, '小说2', '凡骨问长生（全文）.txt')
TARGET = int(sys.argv[2]) if len(sys.argv) > 2 else 1500

def save(state):
    json.dump(state, open(CKPT, 'w', encoding='utf-8'), ensure_ascii=False)

def load():
    if os.path.exists(CKPT):
        try:
            return json.load(open(CKPT, encoding='utf-8'))
        except Exception:
            return {}
    return {}

def llm_chat_long(system, user, max_tokens=4000, timeout=900):
    cfg = ai._load_llm_config()
    errs = []
    for p in cfg:
        body = json.dumps({
            'model': p['model'],
            'messages': [{'role': 'system', 'content': system},
                         {'role': 'user', 'content': user}],
            'max_tokens': max_tokens, 'temperature': 0.3}).encode('utf-8')
        req = urllib.request.Request(p['base'].rstrip('/') + '/chat/completions',
                                     data=body, headers={
                                         'Content-Type': 'application/json',
                                         'Authorization': 'Bearer ' + p['key']})
        try:
            with urllib.request.urlopen(req, timeout=timeout) as r:
                d = json.loads(r.read().decode('utf-8'))
                return d['choices'][0]['message']['content']
        except Exception as e:
            errs.append('%s: %s' % (p['name'], str(e)[:150]))
    raise RuntimeError('LLM 全部不可用: ' + ' | '.join(errs))

def main():
    state = load()
    text = open(TEXT, encoding='utf-8').read()

    # ---- 1) 拆章拆集 ----
    if 'chapters' not in state:
        chapters = ai.split_chapters(text)
        episodes = {}
        for i, ch in enumerate(chapters):
            eps = ai.split_episodes(ch['content'], TARGET)
            episodes[str(i)] = eps
        state['chapters'] = chapters
        state['episodes'] = episodes
        state['phase'] = 'split'
        save(state)
        print('PHASE split chapters=%d episodes=%d' % (len(chapters), sum(len(v) for v in episodes.values())), flush=True)

    # ---- 2) 资产提取（一次大调用，长超时）----
    if state.get('phase') == 'split' or 'assets' not in state:
        chapters = state['chapters']
        parts = []
        for c in chapters:
            take = min(len(c['content']), 40000)
            parts.append(c['title'] + '\n' + c['content'][:take])
        context = '\n\n'.join(parts)[:100000]
        system = ('你是资深网文编辑与漫剧制作人。请通读下面小说节选，'
                  '完整提取制作漫剧所需的结构化信息。只输出一个 JSON 对象，不要任何其他文字。')
        user = ('小说节选（较长，请通读全部内容，不要只看开头）：\n' + context + '\n\n'
                '请输出 JSON，字段如下：\n'
                '{"title":"小说/剧本名",'
                '"characters":[{"name":"角色名","role":"主角/重要配角/反派/路人等","description":"外貌与性格，20字内","outfit":"标志性服装","feature":"标志特征"}],'
                '"scenes":["场景名，如：林府大门·清晨"],'
                '"props":["关键道具名，如：木剑"]}\n'
                '要求：所有出场过且有名字/有戏份的角色都要列入（上限25）；主要地点都列入（上限25）；反复出现或推动剧情的关键道具都列入（上限20）；'
                '角色名必须从文中出现的人名提取，不要编造。')
        raw = llm_chat_long(system, user, max_tokens=4000, timeout=900)
        d = ai._extract_json(raw) or {}
        d.setdefault('characters', [])
        d.setdefault('scenes', [])
        d.setdefault('props', [])
        d.setdefault('title', '未命名剧本')
        fixed = []
        for c in d['characters']:
            if isinstance(c, dict) and c.get('name'):
                fixed.append({'name': str(c['name']).strip(), 'role': str(c.get('role','')).strip(),
                              'description': str(c.get('description','')).strip(),
                              'outfit': str(c.get('outfit','')).strip(), 'feature': str(c.get('feature','')).strip()})
            elif isinstance(c, str) and c.strip():
                fixed.append({'name': c.strip(), 'role': '', 'description': '', 'outfit': '', 'feature': ''})
        d['characters'] = fixed
        state['assets'] = d
        state['phase'] = 'assets'
        save(state)
        print('PHASE assets title=%s chars=%d scenes=%d props=%d' % (
            d.get('title'), len(d['characters']), len(d['scenes']), len(d['props'])), flush=True)

    # ---- 3) 逐集分镜（断点续跑）----
    providers = ai._load_llm_config()
    if 'board' not in state:
        state['board'] = {}
        state['phase'] = 'board'
    chapters = state['chapters']
    episodes = state['episodes']
    total = sum(len(v) for v in episodes.values())
    done = 0
    for i in range(len(chapters)):
        for j in range(len(episodes[str(i)])):
            key = '%d-%d' % (i, j)
            if key in state['board']:
                done += 1
                continue
            try:
                shots = ai.storyboard_episode(episodes[str(i)][j], chapters[i]['title'], j + 1, state['assets'], providers)
            except Exception as e:
                shots = []
                print('ERR %s: %s' % (key, str(e)[:120]), flush=True)
            state['board'][key] = shots
            done += 1
            save(state)
            if done % 3 == 0 or done == total:
                print('PHASE board %d/%d' % (done, total), flush=True)
    state['phase'] = 'done'

    # ---- 4) 组画布 JSON ----
    board = {(int(k.split('-')[0]), int(k.split('-')[1])): v for k, v in state['board'].items()}
    ep_map = {int(k): v for k, v in episodes.items()}
    canvas = ai.build_canvas_json(state['assets'].get('title', '未命名剧本'), chapters, ep_map, state['assets'], board)
    json.dump({'cards': canvas['cards'], 'view': canvas['view']}, open(OUT, 'w', encoding='utf-8'), ensure_ascii=False, indent=2)
    print('PHASE done -> %s' % OUT, flush=True)
    print('DONE', flush=True)

if __name__ == '__main__':
    main()