#!/usr/bin/env python3
"""
DataFlow PDF2Model CLI Module - dataflow/cli_funcs/cli_pdf.py
PDF to Model training pipeline with init/train/chat commands
"""

import copy
import subprocess
import sys
import yaml
import re
import json
import os
import datetime
from pathlib import Path
from colorama import Fore, Style
from dataflow import get_logger
from .paths import DataFlowPath
from transformers import AutoConfig

logger = get_logger()

PDF2MODEL_TRAIN_BACKEND_BASE = "base"

# 可供后续可能的dataflex系列后端接入
PDF2MODEL_TRAIN_BACKEND_DATAFLEX_LESS = "dataflex-less"

DATAFLEX_BACKEND_SPECS = {
    PDF2MODEL_TRAIN_BACKEND_DATAFLEX_LESS: {
        "component_name": "less",
        "suffix": "less",
    },
}


def _get_dataflex_backend_spec(pdf2model_train_backend: str):
    return DATAFLEX_BACKEND_SPECS.get(pdf2model_train_backend)


def _assert_pdf2model_train_backend_allowed(qa_type: str, pdf2model_train_backend: str) -> None:
    if pdf2model_train_backend.startswith("dataflex-"):
        if qa_type != "kbc":
            raise AssertionError(
                f"pdf2model_train_backend={pdf2model_train_backend} is only valid when qa=kbc, got qa={qa_type!r}"
            )
        if _get_dataflex_backend_spec(pdf2model_train_backend) is None:
            supported_backends = ", ".join(sorted(DATAFLEX_BACKEND_SPECS.keys()))
            raise AssertionError(
                f"Unsupported pdf2model_train_backend={pdf2model_train_backend!r}. "
                f"Supported dataflex backends: {supported_backends or '(none)'}"
            )
    if qa_type == "vqa" and pdf2model_train_backend != PDF2MODEL_TRAIN_BACKEND_BASE:
        raise AssertionError(f"vqa only supports pdf2model_train_backend={PDF2MODEL_TRAIN_BACKEND_BASE!r}, got {pdf2model_train_backend!r}")


def _path_for_yaml(cwd: Path, target: Path) -> str:
    """Prefer path relative to cwd for yaml portability."""
    try:
        return str(target.resolve().relative_to(cwd.resolve()))
    except ValueError:
        return str(target.resolve())


def _build_dataflex_components_yaml_text(selector_cache_dir: str, component_name: str, backend_name: str) -> str:
    return f"""# Auto-generated by dataflow pdf2model init ({backend_name}).
selectors:
  {component_name}:
    name: {component_name}
    params:
      cache_dir: {selector_cache_dir}
      gradient_type: adam
      proj_dim: 4096
      seed: 123
      save_interval: 16
"""


def _build_pdf2model_dataflex_config_dict(
    *,
    cwd: Path,
    model_name_or_path: str,
    dataset_train: str,
    dataset_eval: str,
    dataset_dir: Path,
    output_dir: Path,
    components_cfg_file: Path,
    train_config: dict,
    component_name: str,
) -> dict:
    """LlamaFactory-compatible fields + DataFlex dynamic_train block; ``template`` follows ``train_config`` (from LlamaFactoryTrainer defaults)."""
    ds_dir_s = _path_for_yaml(cwd, dataset_dir)
    out_s = _path_for_yaml(cwd, output_dir)
    comp_s = _path_for_yaml(cwd, components_cfg_file)
    # 参数配置，不直接调用dataflex的部件，而是直接构建LlamaFactory的配置，DataFlex的LESS Trainer会从中解析出自己的配置项
    return {
        "model_name_or_path": model_name_or_path,
        "trust_remote_code": train_config.get("trust_remote_code", True),
        "stage": "sft",
        "do_train": True,
        "finetuning_type": train_config.get("finetuning_type", "lora"),
        "lora_target": train_config.get("lora_target", "all"),
        "lora_rank": train_config.get("lora_rank", 16),
        "lora_alpha": train_config.get("lora_alpha", 32),
        "dataset": dataset_train,
        "eval_dataset": dataset_eval,
        "template": train_config.get("template", "qwen"),
        "dataset_dir": ds_dir_s,
        "cutoff_len": train_config.get("cutoff_len", 1024),
        "overwrite_cache": train_config.get("overwrite_cache", True),
        "preprocessing_num_workers": train_config.get("preprocessing_num_workers", 4),
        "dataloader_num_workers": train_config.get("dataloader_num_workers", 0),
        "seed": train_config.get("seed", 42),
        "output_dir": out_s,
        "logging_steps": train_config.get("logging_steps", 10),
        "save_steps": train_config.get("save_steps", 300),
        "plot_loss": train_config.get("plot_loss", True),
        "save_only_model": train_config.get("save_only_model", False),
        "overwrite_output_dir": True,
        "report_to": train_config.get("report_to", "none"),
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": 1,
        "learning_rate": train_config.get("learning_rate", 1.0e-4),
        "num_train_epochs": train_config.get("num_train_epochs", 1.0),
        "lr_scheduler_type": train_config.get("lr_scheduler_type", "cosine"),
        "warmup_ratio": train_config.get("warmup_ratio", 0.1),
        "bf16": train_config.get("bf16", True),
        "fp16": train_config.get("fp16", False),
        "ddp_timeout": train_config.get("ddp_timeout", 180000000),
        "per_device_eval_batch_size": 1,
        "metric_for_best_model": "eval_loss",
        "greater_is_better": False,
        "load_best_model_at_end": False,
        "eval_strategy": "steps",
        "eval_steps": 10,
        "train_type": "dynamic_select",
        "components_cfg_file": comp_s,
        "component_name": component_name,
        "warmup_step": 1,
        "update_step": 1,
        "update_times": 2,
    }


def run_script_with_args(script_path: Path, description: str, args: list = None, cwd: str = None) -> bool:
    """Run a Python script with arguments and real-time output"""
    print(f"\n{Fore.BLUE}{description}{Style.RESET_ALL}")
    cmd = [sys.executable, str(script_path)]
    if args:
        cmd.extend(args)
    print(f"Running: {' '.join(cmd)}")
    if cwd:
        print(f"Working directory: {cwd}")

    try:
        result = subprocess.run(cmd, cwd=cwd, check=True,
                                stdout=sys.stdout, stderr=sys.stderr, text=True)
        print(f"{Fore.GREEN}✅ {description} completed{Style.RESET_ALL}")
        return True
    except subprocess.CalledProcessError as e:
        print(f"{Fore.RED}❌ {description} failed{Style.RESET_ALL}")
        return False


def get_dataflow_script_path(script_name: str) -> Path:
    """Get the path of dataflow built-in scripts"""
    try:
        import dataflow
        dataflow_path = Path(dataflow.__file__).parent

        pdf2model_supportive_kbc_path = dataflow_path / "statics" / "pipelines" / "gpu_pipelines" / "kbcleaning" / script_name
        if pdf2model_supportive_kbc_path.exists():
            return pdf2model_supportive_kbc_path
        
        pdf2model_supportive_vqa_path = dataflow_path / "statics" / "pipelines" / "api_pipelines" / script_name
        if pdf2model_supportive_vqa_path.exists():
            return pdf2model_supportive_vqa_path

        # 检查其他可能的路径
        possible_dirs = [
            dataflow_path / "templates" / "pdf2model_pipeline",
            dataflow_path / "pipeline_templates",
            dataflow_path / "cli_funcs" / "pdf2model_pipeline",
            "./"
        ]

        for dir_path in possible_dirs:
            script_path = dir_path / script_name
            if script_path.exists():
                return script_path

        return None
    except Exception:
        return None


def copy_customizable_scripts(qa_type: str) -> bool:
    """Only copy scripts that users might want to customize"""
    print("Step 0: Copying customizable pipeline script...")

    current_dir = Path(os.getcwd())
    target_name = "pdf_to_model_pipeline.py"

    try:
        if qa_type == "vqa":
            scripts_to_copy = ["pdf_vqa_extract_pipeline.py"]
        else:  # 默认是kbc
            scripts_to_copy = ["kbcleaning_pipeline_vllm.py"]

        import shutil
        copied_files = []

        for script_name in scripts_to_copy:
            source_path = get_dataflow_script_path(script_name)
            if source_path is None:
                print(f"Warning: Template not found: {script_name}")
                continue

            target_path = current_dir / target_name

            shutil.copy2(source_path, target_path)
            copied_files.append(script_name)
            print(f"Copied: {script_name}")

        if copied_files:
            print(f"Successfully copied {len(copied_files)} customizable script(s)")
            return True
        else:
            print("No customizable scripts were copied")
            return False

    except Exception as e:
        print(f"Failed to copy scripts: {e}")
        return False


def create_train_config_yaml(
    cache_path="./",
    model_name_or_path="Qwen/Qwen2.5-7B-Instruct",
    qa_type="kbc",
    pdf2model_train_backend: str = PDF2MODEL_TRAIN_BACKEND_BASE,
):
    """Create train_config.yaml file using built-in LlamaFactory configuration"""
    _assert_pdf2model_train_backend_allowed(qa_type, pdf2model_train_backend)

    cache_path_obj = Path(cache_path)
    if not cache_path_obj.is_absolute():
        caller_cwd = Path(os.environ.get('PWD', os.getcwd()))
        cache_path_obj = caller_cwd / cache_path_obj

    cwd = Path(os.getcwd())

    # 生成时间戳
    timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    model_dir_name = f"pdf2model_cache_{timestamp}"  # 改为pdf2model_cache前缀
    dataflex_spec = _get_dataflex_backend_spec(pdf2model_train_backend)
    if dataflex_spec:
        model_dir_name_dataflex = f"pdf2model_cache_{timestamp}_{dataflex_spec['suffix']}"
    else:
        model_dir_name_dataflex = model_dir_name

    cache_dir = cache_path_obj / ".cache"
    cache_dir.mkdir(parents=True, exist_ok=True)
    config_file = cache_dir / "train_config.yaml"

    try:
        # 使用内置的 LlamaFactory.py 获取默认配置
        llamafactory_script_path = get_dataflow_script_path("llama_factory_trainer.py")
        if llamafactory_script_path is None:
            print("Built-in llama_factory_trainer.py not found")
            return None

        import importlib.util
        spec = importlib.util.spec_from_file_location("llamafactory_trainer", llamafactory_script_path)
        llamafactory_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(llamafactory_module)

        # 创建trainer实例并获取默认配置
        trainer = llamafactory_module.LlamaFactoryTrainer(str(config_file), str(cache_path_obj))
        config = trainer.get_default_config()
        # 只更新必要的动态参数
        if model_name_or_path:
            config["model_name_or_path"] = model_name_or_path
        config["output_dir"] = str(cache_path_obj / ".cache" / "saves" / model_dir_name)
        config["dataset_dir"] = str(cache_path_obj / ".cache" / "data")
        if qa_type == "vqa":
            dataset_name = "pdf_vqa_dataset"
        else:
            dataset_name = "pdf_kbc_dataset"
        config["dataset"] = dataset_name

        dataflex_yaml_path = None
        components_yaml_path = cache_dir / "pdf2model_dataflex_components.yaml"
        if dataflex_spec:
            dataflex_yaml_path = cache_dir / f"pdf2model_dataflex_{dataflex_spec['suffix']}.yaml"
            selector_cache_rel = _path_for_yaml(
                cwd, cache_dir / f"dataflex_{dataflex_spec['suffix']}_selector_cache"
            )

        if dataflex_spec:
            with open(components_yaml_path, "w", encoding="utf-8") as cf:
                cf.write(
                    _build_dataflex_components_yaml_text(
                        selector_cache_rel,
                        component_name=dataflex_spec["component_name"],
                        backend_name=pdf2model_train_backend,
                    )
                )

            dataset_eval_name = "pdf_kbc_eval_dataset"
            dataflex_config = _build_pdf2model_dataflex_config_dict(
                cwd=cwd,
                model_name_or_path=config["model_name_or_path"],
                dataset_train=dataset_name,
                dataset_eval=dataset_eval_name,
                dataset_dir=cache_path_obj / ".cache" / "data",
                output_dir=cache_path_obj / ".cache" / "saves" / model_dir_name_dataflex,
                components_cfg_file=components_yaml_path,
                train_config=config,
                component_name=dataflex_spec["component_name"],
            )
            with open(dataflex_yaml_path, "w", encoding="utf-8") as lf:
                yaml.dump(dataflex_config, lf, default_flow_style=False, allow_unicode=True, sort_keys=False, indent=2)
            print(f"{dataflex_yaml_path.name} created ({pdf2model_train_backend}): {dataflex_yaml_path}")
            print(f"pdf2model_dataflex_components.yaml created: {components_yaml_path}")

        pdf2model_state = {
            "qa": qa_type,
            "pdf2model_train_backend": pdf2model_train_backend,
            "train_config_file_dir": str(config_file),
            "output_dir": str(cache_path_obj / ".cache" / "saves" / model_dir_name),
            "dataset": dataset_name,
            "dataset_dir": str(cache_path_obj / ".cache" / "data"),
            "timestamp": timestamp,
        }
        if dataflex_spec:
            pdf2model_state["dataflex_backend_component"] = dataflex_spec["component_name"]
            pdf2model_state["dataflex_train_yaml"] = str(dataflex_yaml_path)
            pdf2model_state["dataflex_output_dir"] = str(
                cache_path_obj / ".cache" / "saves" / model_dir_name_dataflex
            )
            pdf2model_state["dataflex_components_yaml"] = str(components_yaml_path)

        with open(cache_dir / "pdf2model_state.json", 'w') as f:
            json.dump(pdf2model_state, f, indent=2)

        # 保存配置
        with open(config_file, 'w', encoding='utf-8') as f:
            yaml.dump(config, f,
                      default_flow_style=False,
                      allow_unicode=True,
                      sort_keys=False,
                      indent=2)

        print(f"train_config.yaml created: {config_file}")
        print(f"Model will be saved to: {model_dir_name}")
        return str(config_file)

    except Exception as e:
        print(f"Failed to create train_config.yaml: {e}")
        return None

def generate_dataset_info(
    cache_path_obj,
    dataset_name,
    qa_type,
    pdf2model_train_backend: str = PDF2MODEL_TRAIN_BACKEND_BASE,
):
    """Create dataset_info.json configuration automatically"""
    _assert_pdf2model_train_backend_allowed(qa_type, pdf2model_train_backend)

    dataset_info_path = cache_path_obj / ".cache" / "data" / "dataset_info.json"
    dataset_info_path.parent.mkdir(parents=True, exist_ok=True)

    if qa_type == "vqa":
        config_entry = {
            "file_name": "qa.json",
            "formatting": "sharegpt",
            "columns": {
                "messages": "messages", 
                "images": "images"       
            },
            "tags": {
                "role_tag": "role",
                "content_tag": "content",
                "user_tag": "user",
                "assistant_tag": "assistant"
            }            
        }
    else:
        config_entry = {
            "file_name": "qa.json",
            "formatting": "alpaca",
            "columns": {
                "prompt": "instruction",
                "query": "input",
                "response": "output"
            }
        }

    # duplicate entry for eval when using DataFlex backend (same file; backend needs eval_dataset registered)
    if (
        qa_type == "kbc"
        and _get_dataflex_backend_spec(pdf2model_train_backend) is not None
        and dataset_name == "pdf_kbc_dataset"
    ):
        dataset_info = {
            dataset_name: copy.deepcopy(config_entry),
            "pdf_kbc_eval_dataset": copy.deepcopy(config_entry),
        }
    else:
        dataset_info = {dataset_name: config_entry}

    try:
        with open(dataset_info_path, 'w', encoding='utf-8') as f:
            json.dump(dataset_info, f, indent=2, ensure_ascii=False)
        print(f"✅ Dataset info registered as '{dataset_name}' with {qa_type} format.")
        return True
    except Exception as e:
        print(f"❌ Failed to write dataset_info.json: {e}")
        return False

def verify_environment():
    """Verify runtime environment"""
    print("Checking environment...")

    missing_deps = []

    try:
        import llamafactory
        print("✅ LlamaFactory installed")
    except ImportError:
        missing_deps.append("llamafactory[torch,metrics]")

    try:
        import yaml
        print("✅ PyYAML installed")
    except ImportError:
        missing_deps.append("pyyaml")

    if missing_deps:
        print(f"❌ Missing dependencies: {', '.join(missing_deps)}")
        print(f"Install with: pip install {' '.join(missing_deps)}")
        return False

    return True


def check_required_files():
    """Check if required built-in scripts exist"""
    # 检查所有需要的内置脚本
    required_scripts = [
        "path_to_jsonl_script.py",
        "llama_factory_trainer.py"
    ]

    missing_scripts = []
    for script in required_scripts:
        script_path = get_dataflow_script_path(script)
        if script_path is None:
            missing_scripts.append(script)
        else:
            print(f"✅ Found built-in script: {script}")

    if missing_scripts:
        print(f"❌ Missing built-in scripts: {', '.join(missing_scripts)}")
        print("These should be part of the dataflow installation")
        return False

    # 检查用户目录下是否有可自定义的脚本
    current_dir = Path(os.getcwd())
    customizable_script = current_dir / "pdf_to_model_pipeline.py"
    if customizable_script.exists():
        print("✅ Found customizable script: pdf_to_model_pipeline.py")
    else:
        print("❌ Missing customizable script: pdf_to_model_pipeline.py")
        print("Run 'dataflow pdf2model init' first")
        return False

    return True


def cli_pdf2model_init(
    cache_path: str = "./",
    model_name: str = None,
    qa_type: str = "kbc",
    pdf2model_train_backend: str = PDF2MODEL_TRAIN_BACKEND_BASE,
) -> bool:
    """
    PDF2Model initialization:
    0. Copy only customizable scripts to current directory
    1. Create train_config.yaml in .cache directory
    """
    _assert_pdf2model_train_backend_allowed(qa_type, pdf2model_train_backend)

    effective_model = model_name if model_name else "Qwen/Qwen2.5-7B-Instruct"
    print("Starting PDF2Model initialization...")
    print(f"Cache directory: {cache_path}")
    print(f"Model: {effective_model}")
    print(f"Train backend: {pdf2model_train_backend}")
    print(f"Output directory: pdf2model_cache_<timestamp>")  # 更新输出目录显示
    print("-" * 60)

    if not verify_environment():
        return False

    try:
        # Step 0: Copy only customizable scripts
        if not copy_customizable_scripts(qa_type):
            return False

        # Step 1: Create training configuration
        print("Step 1: Creating training configuration...")
        config_file = create_train_config_yaml(
            cache_path,
            effective_model,
            qa_type,
            pdf2model_train_backend=pdf2model_train_backend,
        )

        if config_file:
            print("PDF2Model initialization completed!")
            return True
        else:
            print("Failed to create training configuration")
            return False

    except Exception as e:
        print(f"Initialization failed: {e}")
        return False


def get_latest_model_dir(cache_path_obj):
    """Latest adapter directory under saves (by mtime; supports pdf2model_cache_* suffixes e.g. _less)."""
    saves_dir = cache_path_obj / ".cache" / "saves"
    if not saves_dir.exists():
        return None

    candidates = []
    for dir_path in saves_dir.iterdir():
        if dir_path.is_dir() and dir_path.name.startswith("pdf2model_cache_"):
            candidates.append((dir_path.stat().st_mtime, dir_path))

    if not candidates:
        return None

    candidates.sort(key=lambda x: x[0], reverse=True)
    return candidates[0][1]


def cli_pdf2model_train(lf_yaml: str = ".cache/train_config.yaml", cache_path: str = "./") -> bool:
    """
    Start PDF2Model training using mix of built-in and user scripts
    """
    print("Starting PDF2Model training...")

    current_dir = Path(os.getcwd())

    cache_path_obj = Path(cache_path)
    if not cache_path_obj.is_absolute():
        cache_path_obj = current_dir / cache_path_obj

    train_cfg_path = cache_path_obj / ".cache" / "train_config.yaml"
    if not train_cfg_path.exists():
        print(f"Training config file not found: {train_cfg_path}")
        print(f"{Style.BRIGHT}Run 'dataflow pdf2model init' first")
        return False

    try:
        with open(train_cfg_path, "r", encoding="utf-8") as f:
            train_config = yaml.safe_load(f) or {}
    except Exception as e:
        print(f"[ERROR] Could not read train_config.yaml: {e}")
        return False

    pdf2model_state_path = cache_path_obj / ".cache" / "pdf2model_state.json"
    try:
        with open(pdf2model_state_path, "r", encoding="utf-8") as f:
            pdf2model_state = json.load(f)
            print(f"Loaded pipeline state from {pdf2model_state_path}")
    except Exception as e:
        print(f"Warning: State file not found at {pdf2model_state_path}, using defaults.")
        pdf2model_state = {}

    qa_type = pdf2model_state.get("qa")
    pdf2model_train_backend = pdf2model_state.get(
        "pdf2model_train_backend", PDF2MODEL_TRAIN_BACKEND_BASE
    )

    try:
        _assert_pdf2model_train_backend_allowed(qa_type or "kbc", pdf2model_train_backend)
    except AssertionError as err:
        print(f"{Fore.RED}Invalid pdf2model state: {err}{Style.RESET_ALL}")
        return False

    if not verify_environment():
        return False

    if _get_dataflex_backend_spec(pdf2model_train_backend):
        from dataflow.cli_funcs.pdf2model_pipeline.dataflex_pdf2model_launcher import verify_dataflex_available

        if not verify_dataflex_available():
            return False

    if not check_required_files():
        return False

    lf_path = Path(lf_yaml)
    if not lf_path.is_absolute():
        lf_path = current_dir / lf_path
    if pdf2model_train_backend == PDF2MODEL_TRAIN_BACKEND_BASE and not lf_path.exists():
        print(f"LlamaFactory training config not found: {lf_path}")
        print(f"{Style.BRIGHT}Run 'dataflow pdf2model init' first or pass a valid --lf-yaml")
        return False

    print("-" * 60)

    try:
        # Step 1: PDF Detection
        script1_path = get_dataflow_script_path("path_to_jsonl_script.py")
        args1 = ["./", "--output", str(cache_path_obj / ".cache" / "pdf_list.jsonl"), "--qa_type", str(qa_type)]
        if not run_script_with_args(script1_path, "Step 1: PDF Detection", args1, cwd=str(current_dir)):
            return False

        # Step 2: Data Processing
        target_script_name = "pdf_to_model_pipeline.py"
        script2 = current_dir / target_script_name
        args2 = ["--cache", cache_path]
        if not run_script_with_args(script2, "Step 2: Data Processing", args2, cwd=str(current_dir)):
            return False

        # Step 2.5: Create dataset_info.json (dynamically)
        print(f"\n{Fore.BLUE}Step 2.5: Creating dataset_info.json{Style.RESET_ALL}")

        # 获取数据集名称
        dataset_name = train_config.get("dataset")
        if isinstance(dataset_name, list):
            dataset_name = dataset_name[0]  # 如果是列表，取第一个

        if not dataset_name:
            print("Warning: No dataset name found in train_config.yaml, using default 'kb_qa'")
            dataset_name = "kb_qa"

        print(f"Dataset name from config: {dataset_name}")

        if not generate_dataset_info(
            cache_path_obj,
            dataset_name,
            qa_type,
            pdf2model_train_backend=pdf2model_train_backend,
        ):
            return False

        print(f"{Fore.GREEN}✅ Step 2.5: Creating dataset_info.json completed{Style.RESET_ALL}")

        # Step 3: Data Conversion - skip
        print(f"\n{Fore.BLUE}Step 3: Data Conversion{Style.RESET_ALL}")
        qa_json_path = cache_path_obj / ".cache" / "data" / "qa.json"
        if qa_json_path.exists():
            print(f"✅ qa.json already in correct format, skipping conversion")
            print(f"{Fore.GREEN}✅ Step 3: Data Conversion completed{Style.RESET_ALL}")
        else:
            print(f"❌ qa.json not found at {qa_json_path}")
            return False

        # Step 4: Training
        if pdf2model_train_backend == PDF2MODEL_TRAIN_BACKEND_BASE:
            script4_path = get_dataflow_script_path("llama_factory_trainer.py")
            args4 = ["--config", str(lf_path), "--cache", cache_path]
            if not run_script_with_args(script4_path, "Step 4: Training (LlamaFactory)", args4, cwd=str(current_dir)):
                return False
            try:
                with open(lf_path, "r", encoding="utf-8") as f:
                    config = yaml.safe_load(f)
                    actual_output_dir = config.get("output_dir", "unknown")
            except Exception:
                actual_output_dir = "unknown"
        else:
            from dataflow.cli_funcs.pdf2model_pipeline.dataflex_pdf2model_launcher import run_dataflex_train

            dfx_yaml = pdf2model_state.get("dataflex_train_yaml")
            if not dfx_yaml:
                print(f"{Fore.RED}Missing dataflex_train_yaml in pdf2model_state.json; re-run pdf2model init.{Style.RESET_ALL}")
                return False
            dfx_path = Path(dfx_yaml)
            if not dfx_path.is_absolute():
                dfx_path = current_dir / dfx_path
            if not dfx_path.exists():
                print(f"{Fore.RED}DataFlex train yaml not found: {dfx_path}{Style.RESET_ALL}")
                return False
            if not run_dataflex_train(dfx_path, current_dir):
                return False
            actual_output_dir = pdf2model_state.get("dataflex_output_dir", "unknown")

        print("Training completed successfully!")
        print(f"Model saved to: {actual_output_dir}")
        print("Next steps:")
        print(f"{Style.BRIGHT}Test the trained model with 'dataflow chat'")

        return True

    except Exception as e:
        print(f"Training error: {e}")
        return False


def cli_pdf2model_chat(model_path=None, cache_path="./", base_model=None):
    """Start LlamaFactory chat interface"""
    current_dir = Path(os.getcwd())

    # 处理cache路径
    cache_path_obj = Path(cache_path)
    if not cache_path_obj.is_absolute():
        cache_path_obj = current_dir / cache_path_obj

    # 确定模型路径
    if model_path is None:
        # 获取最新的模型目录
        latest_model_dir = get_latest_model_dir(cache_path_obj)
        if latest_model_dir:
            model_path = latest_model_dir
        else:
            print("No trained model found")
            print("Run 'dataflow pdf2model train' to train a model first")
            return False

    model_path = Path(model_path)
    if not model_path.exists():
        print(f"Model not found: {model_path}")
        print("Run 'dataflow pdf2model train' to train a model first")
        return False

    # 验证是否为有效的adapter目录
    adapter_files = [
        "adapter_config.json",
        "adapter_model.bin",
        "adapter_model.safetensors"
    ]

    has_adapter = any((model_path / f).exists() for f in adapter_files)
    if not has_adapter:
        print(f"No adapter files found in {model_path}")
        print("This doesn't appear to be a trained adapter directory.")
        print("Expected files: adapter_config.json, adapter_model.bin/safetensors")
        return False

    # 确定基础模型路径 - 安全的读取方式
    if base_model is None:
        base_model = None  # 先设为None

        # 尝试从训练配置中读取基础模型
        config_file = cache_path_obj / ".cache" / "train_config.yaml"
        if config_file.exists():
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    config = yaml.safe_load(f)
                    base_model = config.get('model_name_or_path')
                    if base_model:
                        print(f"Found base model in config: {base_model}")
            except Exception as e:
                print(f"Warning: Could not read config file: {e}")

        # 尝试从adapter_config.json读取
        if not base_model:
            adapter_config_path = model_path / "adapter_config.json"
            if adapter_config_path.exists():
                try:
                    with open(adapter_config_path, 'r', encoding='utf-8') as f:
                        adapter_config = json.load(f)
                        base_model = adapter_config.get('base_model_name_or_path')
                        if base_model:
                            print(f"Found base model in adapter config: {base_model}")
                except Exception as e:
                    print(f"Warning: Could not read adapter config: {e}")

        # 如果仍然没有找到base_model，报错退出而不是使用默认值
        if not base_model:
            print("Cannot determine base model path")
            print("Please ensure your training config contains 'model_name_or_path'")
            print("Or check that adapter_config.json exists and contains 'base_model_name_or_path'")
            return False

    # 检查LlamaFactory
    try:
        import llamafactory
        print("LlamaFactory available")
    except ImportError:
        print("LlamaFactory not installed")
        print("Install with: pip install llamafactory[torch,metrics]")
        return False

    # 直接用命令行参数启动聊天
    chat_cmd = [
        "llamafactory-cli", "chat",
        "--model_name_or_path", base_model,
        "--adapter_name_or_path", str(model_path.absolute())
    ]

    print(f"Base model: {base_model}")
    print(f"Adapter path: {model_path}")
    print(f"Command: {' '.join(chat_cmd)}")
    print("-" * 60)
    print("Starting chat session...")
    print("-" * 60)

    try:
        result = subprocess.run(chat_cmd, check=True)
        print("\nChat session completed")
        return True
    except subprocess.CalledProcessError as e:
        print(f"\nChat failed: {e}")
        return False
    except KeyboardInterrupt:
        print("\n\nChat session ended by user")
        return True