import os, json, pandas as pd, csv
from typing import Dict, Any, List, Set
from config import VLLMConfig
from utils import norm_key, ensure_dir
from sampling import read_video_meta, sample_indices_uniform, sample_indices_fps, grab_frames
from client import to_jpeg_base64, call_vllm
from prompt import build_messages
from parser import parse_step5_line, validate_bbox_sequences
from video_utils import compose_videos, build_out_paths
from tqdm import tqdm

def load_dataframe(csv_path: str) -> pd.DataFrame:
    # 同前版本：兼容你提供的表头
    with open(csv_path, 'r', encoding='utf-8') as f:
        reader = csv.reader(f)
        try:
            header = next(reader)
        except StopIteration:
            return pd.DataFrame(columns=['videoid','source_video_path','instruction','chinese_instruction'])
        norm_header = [h.strip() for h in header]
        if len(norm_header) >= 4 and \
           norm_header[0].lower() == 'videoid' and \
           norm_header[1].lower() == 'source_video_path' and \
           norm_header[2].lower() == 'instruction' and \
           norm_header[3].lower() == 'chinese_instruction':
            rows = []
            for row in reader:
                if not row or all((x is None or str(x).strip() == '') for x in row): continue
                vid  = (row[0].strip() if len(row) > 0 else '')
                vpath= (row[1].strip() if len(row) > 1 else '')
                rest = row[2:] if len(row) > 2 else []
                joined = ','.join([x if x is not None else '' for x in rest])
                instr, zh = ('','')
                if ',' in joined: instr, zh = joined.rsplit(',', 1)
                else: instr, zh = joined, ''
                rows.append({
                    'videoid': vid,
                    'source_video_path': vpath,
                    'instruction': instr.strip(),
                    'chinese_instruction': zh.strip(),
                })
            return pd.DataFrame(rows, columns=['videoid','source_video_path','instruction','chinese_instruction'])
        else:
            return pd.read_csv(csv_path, engine='python')

def _stable_key(row: Dict[str,Any]) -> str:
    cols = {norm_key(k): k for k in row.keys()}
    def pick(*cands, default=None):
        for c in cands:
            if c in cols: return row[cols[c]]
        return default
    uid = pick('videoid','uid','id','task_id','clip_id','name','key', default=None)
    video_path = pick('source_video_path','video_path','video','filepath','file_path')
    sampling_mode = pick('sampling_mode','sampling','mode', default='uniform_5')
    stem = str(uid) if uid not in (None, "") else os.path.splitext(os.path.basename(video_path or ""))[0]
    return f"{stem}::{sampling_mode}"

def _append_results_csv(results_csv: str, row: Dict[str,Any], header_written: bool):
    df = pd.DataFrame([row])
    write_header = not os.path.exists(results_csv) or (os.path.getsize(results_csv)==0 and not header_written)
    df.to_csv(results_csv, mode='a', header=write_header, index=False, encoding='utf-8-sig')

def process_one(row: Dict[str,Any], vllm: VLLMConfig, outroot: str, dry_run: bool=False) -> Dict[str,Any]:
    cols = {norm_key(k): k for k in row.keys()}
    def pick(*cands, default=None):
        for c in cands:
            if c in cols: return row[cols[c]]
        return default

    video_path = pick('source_video_path','video_path','video','filepath','file_path')
    cn = pick('chinese_instruction'); en = pick('instruction','edit_instruction','edit','prompt','text')
    edit_instruction = (cn if cn and str(cn).strip() else en)
    sampling_mode = pick('sampling_mode','sampling','mode', default='uniform_5')
    fps_val = pick('fps', default=None); n_frames_val = pick('n_frames','n','nframes','num_frames', default=None)
    uid = pick('videoid','uid','id','task_id','clip_id','name','key', default=None)

    fps = 1.0 if fps_val in (None,"") else float(fps_val)
    n_frames = 5 if n_frames_val in (None,"") else int(n_frames_val)

    if not video_path or not os.path.exists(video_path):
        raise AssertionError(f'视频不存在或未提供：{video_path}')
    if not (isinstance(edit_instruction,str) and len(str(edit_instruction).strip())>0):
        raise AssertionError('缺少编辑指令（chinese_instruction / instruction）')

    meta = read_video_meta(video_path)
    if sampling_mode == 'fps_1':
        indices = sample_indices_fps(meta['total'], meta['fps'], fps or 1.0)
        mode_tag = f'fps_{int(fps)}'
    else:
        indices = sample_indices_uniform(meta['total'], n_frames or 5)
        mode_tag = f'uniform_{len(indices)}'
    frames_rgb, mapping = grab_frames(video_path, indices)
    frames_b64 = [to_jpeg_base64(f) for f in frames_rgb]
    messages = build_messages(frames_b64, edit_instruction)

    step5_line = ""
    raw_resp = None
    latency = None
    if not dry_run:
        step5_line, latency, raw_resp = call_vllm(vllm.api_base, vllm.model, messages)
        # 记录“原始文本输出”（完整保留，尚未做正则抽取）
        raw_dir = os.path.join(outroot, 'raw_text')
        os.makedirs(raw_dir, exist_ok=True)
        stem = str(uid) if uid not in (None, "") else os.path.splitext(os.path.basename(video_path))[0]
        raw_txt_path = os.path.join(raw_dir, f'{stem}__{mode_tag}.txt')
        with open(raw_txt_path, 'w', encoding='utf-8') as fw:
            fw.write(step5_line if isinstance(step5_line, str) else str(step5_line))
        # 仅取最后一行作为 Step5 行（若模型多段输出）
        lines = [l.strip() for l in step5_line.strip().splitlines() if l.strip()]
        step5_line = lines[-1] if lines else step5_line
    else:
        raw_txt_path = ""

    parsed = {}
    enhanced_instruction = ""
    tasks = {}
    if step5_line:
        parsed = parse_step5_line(step5_line)  # 按你固定格式解析
        validate_bbox_sequences(parsed)        # 坐标规范校验
        enhanced_instruction = parsed.get('enhanced_instruction','')
        tasks = parsed.get('tasks',{})

    # 输出路径（含原视频名 + 后缀）
    stem = str(uid) if uid not in (None, "") else os.path.splitext(os.path.basename(video_path))[0]
    out_inspect, out_mask = build_out_paths(outroot, video_path, stem, mode_tag)
    out_json = os.path.join(outroot, 'parsed_json', f'{stem}__{mode_tag}.json')
    for d in [os.path.dirname(out_inspect), os.path.dirname(out_mask), os.path.dirname(out_json)]:
        os.makedirs(d, exist_ok=True)

    # 无论是否解析成功，都保存 JSON（含 raw_response）
    with open(out_json, 'w', encoding='utf-8') as f:
        json.dump({
            'video_path': video_path,
            'sampling_mode': sampling_mode,
            'indices': indices,
            'mapping_local_t_to_global_index': mapping,
            'model_api_base': vllm.api_base,
            'model_name': vllm.model,
            'enhanced_instruction': enhanced_instruction,
            'tasks': tasks,
            'raw_response': raw_resp,
            'raw_text_path': raw_txt_path,
            'parsed_ok': bool(parsed),
        }, f, ensure_ascii=False, indent=2)

    # 两路视频：采样帧原视频(红框) + 黑底白框；无坐标也生成，便于排查
    compose_videos(video_path, mapping, parsed or {"tasks": {}}, out_inspect, out_mask, allow_empty=True)

    return {
        'key': f"{stem}::{sampling_mode}",
        'uid': uid,
        'video_path': video_path,
        'edit_instruction_src': edit_instruction,
        'enhanced_instruction': enhanced_instruction,
        'sampling_mode': sampling_mode,
        'n_sampled': len(indices),
        'indices_json': json.dumps(indices, ensure_ascii=False),
        'mapping_json': json.dumps(mapping, ensure_ascii=False),
        'tasks_json': json.dumps(tasks, ensure_ascii=False),
        'latency_s': latency,
        'out_inspect': out_inspect,
        'out_mask': out_mask,
        'out_json': out_json,
        'raw_text_path': raw_txt_path,
        'parsed_ok': bool(parsed),
        'error': '' if parsed else 'parsed_empty_or_invalid'
    }

def run_csv(csv_path: str, outroot: str, vllm: VLLMConfig, dry_run: bool=False,
            results_csv: str=None, resume: bool=False, show_progress: bool=True) -> List[Dict[str,Any]]:
    df = load_dataframe(csv_path)
    done_keys: Set[str] = set()
    header_written = False
    if results_csv and os.path.exists(results_csv) and resume:
        try:
            existing = pd.read_csv(results_csv)
            if 'key' in existing.columns:
                done_keys = set(existing['key'].astype(str).tolist())
        except Exception:
            pass

    iterator = df.iterrows()
    if show_progress:
        iterator = tqdm(df.iterrows(), total=len(df), desc="Processing", ncols=100)

    results = []
    for _, row in iterator:
        rdict = {}
        try:
            key = _stable_key(row.to_dict())
            if resume and key in done_keys:
                continue
            rdict = process_one(row.to_dict(), vllm, outroot, dry_run=dry_run)
        except Exception as e:
            rdict = {
                'key': _stable_key(row.to_dict()),
                'uid': row.to_dict().get('videoid', None),
                'video_path': row.to_dict().get('source_video_path', None),
                'edit_instruction_src': row.to_dict().get('chinese_instruction', '') or row.to_dict().get('instruction', ''),
                'enhanced_instruction': '',
                'sampling_mode': row.to_dict().get('sampling_mode',''),
                'n_sampled': None,
                'indices_json': '',
                'mapping_json': '',
                'tasks_json': '',
                'latency_s': None,
                'out_inspect': '',
                'out_mask': '',
                'out_json': '',
                'raw_text_path': '',
                'parsed_ok': False,
                'error': str(e)
            }
        results.append(rdict)
        if results_csv:
            _append_results_csv(results_csv, rdict, header_written=header_written)
            header_written = True
            # 断点续做日志
            ckpt_path = os.path.join(outroot, 'logs', 'pipeline_checkpoint.jsonl')
            with open(ckpt_path, 'a', encoding='utf-8') as fp:
                fp.write(json.dumps(rdict, ensure_ascii=False) + '\n')
    return results