main.py

#!/usr/bin/env python3
"""
Insight Agent — 智能数据分析爬虫智能体
===========================================
CLI 交互式工具，集成爬虫采集 → 数据分析 → 报告生成 → AI问答

用法:
  python3 main.py           # 交互模式
  insight                    # 全局命令（需安装到 PATH）
"""

import sys
import json
import csv
import logging
from pathlib import Path
from datetime import datetime
from typing import List, Dict

sys.path.insert(0, str(Path(__file__).parent))

from logging_config import setup_logging
setup_logging()

from rich.console import Console
from rich.panel import Panel
from rich.text import Text
from rich.markdown import Markdown
from rich.prompt import Prompt

from prompt_toolkit import PromptSession
from prompt_toolkit.styles import Style
from prompt_toolkit.history import InMemoryHistory
from prompt_toolkit.completion import WordCompleter

from config import (
    PROJECT_ROOT, REPORT_DIR, DATA_DIR, ANALYSIS_MODELS,
    THEME, LLM_API_KEY, LLM_BASE_URL, LLM_MODEL,
    CRAWLER_SCRIPTS,
)
from ui import (
    show_welcome,
    select_mode, select_existing_file, select_multiple_files,
    select_crawler, get_crawler_args, select_analysis_models, confirm_step,
    save_crawler_state, save_mode_state,
)
from crawlers.runner import run_crawler, get_latest_output_files
from analysis import run_quantitative, run_qualitative
from report import generate_report
from chat import AnalystSession
from chat.agent import SESSION_DIR

logger = logging.getLogger(__name__)
console = Console()

PROMPT_STYLE = Style.from_dict({
    "prompt": "bold cyan",
    "separator": "dim",
})


# ================================================================
# 数据加载与预览
# ================================================================

def load_data(file_path: Path) -> list:
    """从文件加载数据"""
    ext = file_path.suffix.lower()

    try:
        if ext == ".csv":
            with open(file_path, "r", encoding="utf-8-sig", errors="replace") as f:
                return list(csv.DictReader(f))
        elif ext == ".xlsx":
            try:
                import openpyxl
                wb = openpyxl.load_workbook(file_path, read_only=True, data_only=True)
                ws = wb.active
                rows = list(ws.iter_rows(values_only=True))
                if not rows:
                    return []
                headers = [str(h) if h else f"col_{i}" for i, h in enumerate(rows[0])]
                data = []
                for row in rows[1:]:
                    data.append({headers[i]: str(v) if v is not None else "" for i, v in enumerate(row)})
                wb.close()
                return data
            except ImportError:
                logger.error("缺少 openpyxl 依赖，请运行: pip install openpyxl")
                console.print("[red]缺少 openpyxl 依赖，请运行: pip install openpyxl[/]")
                return []
            except Exception as e:
                logger.error(f"读取Excel失败: {e}")
                console.print(f"[red]读取Excel失败: {e}[/]")
                return []
        elif ext == ".json":
            with open(file_path, "r", encoding="utf-8") as f:
                content = json.load(f)
                if isinstance(content, list):
                    return content
                elif isinstance(content, dict):
                    for v in content.values():
                        if isinstance(v, list) and v:
                            return v
                    return [content]
    except FileNotFoundError:
        logger.error(f"文件不存在: {file_path}")
        console.print(f"[red]文件不存在: {file_path}[/]")
        return []
    except PermissionError:
        logger.error(f"没有文件读取权限: {file_path}")
        console.print(f"[red]没有文件读取权限: {file_path}[/]")
        return []
    except Exception as e:
        logger.error(f"加载数据失败: {e}")
        console.print(f"[red]加载数据失败: {e}[/]")
        return []

    return []


def preview_data(data: list, data_file: Path) -> bool:
    """预览数据：展示前10行 + 字段类型检测，确认是否继续"""
    from rich.table import Table
    from analysis.quantitative import auto_detect_fields

    console.print()
    console.print(Panel.fit(
        f"[bold cyan]数据预览: {data_file.name}[/]",
        border_style="cyan"
    ))

    file_size = data_file.stat().st_size
    size_str = f"{file_size/1024:.1f}KB" if file_size < 1024*1024 else f"{file_size/1024/1024:.1f}MB"
    console.print(f"  [dim]记录数:[/] [bold]{len(data)}[/]  |  "
                f"[dim]字段数:[/] [bold]{len(data[0]) if data else 0}[/]  |  "
                f"[dim]文件大小:[/] [bold]{size_str}[/]")
    console.print()

    field_types = auto_detect_fields(data)
    type_icons = {"text": "📝", "numeric": "🔢", "date": "📅", "other": "📦"}
    type_colors = {"text": "cyan", "numeric": "green", "date": "yellow", "other": "dim"}

    console.print("  [bold]字段类型检测:[/]")
    type_strs = []
    for f, t in field_types.items():
        icon = type_icons.get(t, "•")
        color = type_colors.get(t, "white")
        type_strs.append(f"[{color}]{icon} {f} ({t})[/]")
    console.print("    " + "  |  ".join(type_strs))
    console.print()

    if data:
        console.print("  [bold]数据预览 (前10行):[/]")
        table = Table(show_header=True, header_style="bold white", border_style="dim",
                     show_lines=False, padding=(0, 1))

        columns = list(data[0].keys())[:8]
        for col in columns:
            table.add_column(col[:12], style="dim", max_width=15, no_wrap=True)

        for i, row in enumerate(data[:10]):
            values = [str(row.get(c, ""))[:15] for c in columns]
            table.add_row(*values)

        console.print(table)

    console.print()
    return confirm_step("确认使用此数据进行分析？")


# ================================================================
# 分析调度
# ================================================================

def _run_single_quant(data, model, data_file):
    """单个定量分析任务（线程安全）"""
    try:
        from analysis.quantitative import basic_stats, sentiment_analysis, word_frequency, time_series
        models_map = {
            "basic_stats": basic_stats,
            "sentiment": sentiment_analysis,
            "word_freq": word_frequency,
            "time_series": time_series,
        }
        fn = models_map.get(model)
        if fn:
            return {model: fn(data, data_file)}
        return {model: {"model": model, "error": f"未知模型: {model}"}}
    except Exception as e:
        logger.error(f"定量分析失败 [{model}]: {e}")
        return {model: {"model": model, "error": f"分析失败: {e}"}}


def _run_single_qual(data, model, data_file):
    """单个定性分析任务（线程安全）"""
    try:
        from analysis.qualitative import theme_extraction, content_category, insight_mining
        models_map = {
            "theme_extraction": theme_extraction,
            "content_category": content_category,
            "insight_mining": insight_mining,
        }
        fn = models_map.get(model)
        if fn:
            return {model: fn(data)}
        return {model: {"model": model, "error": f"未知模型: {model}"}}
    except Exception as e:
        logger.error(f"定性分析失败 [{model}]: {e}")
        return {model: {"model": model, "error": f"分析失败: {e}"}}


# ================================================================
# AI 分析对话
# ================================================================

def run_ai_analysis_flow(data: list, data_file: Path):
    """AI 驱动分析流程：数据预览后 AI 直接主持分析 + 持续对话"""
    console.print()
    console.print(Panel.fit(
        "[bold cyan]🤖 AI 分析师已就绪[/]\n"
        f"数据: [dim]{data_file.name}[/] — {len(data)} 条记录, {len(data[0]) if data else 0} 个字段\n\n"
        "[dim]AI 正在查看你的数据...[/]",
        border_style="cyan"
    ))

    session = AnalystSession(data, data_file)

    console.print("\n  [dim]AI 分析中...[/]\n")
    greeting = session.analyst.greet_and_analyze()
    console.print(Panel(
        Markdown(greeting),
        title="🤖 AI 分析师",
        border_style="cyan",
        padding=(1, 2),
    ))

    chat_loop_analyst(session)


def chat_loop_analyst(session: AnalystSession):
    """AI 分析对话循环"""
    completer = WordCompleter(["情感分析", "词频统计", "趋势分析", "总结", "建议", "exit", "help"])
    console.print(Panel.fit(
        "[bold]现在可以指挥 AI 做分析了[/]\n"
        "[dim]输入分析需求，AI 自动写代码执行 | exit 退出 | help 帮助[/]",
        border_style="dim"
    ))

    prompt_session = PromptSession(
        history=InMemoryHistory(),
        style=PROMPT_STYLE,
        completer=completer,
    )

    while True:
        try:
            user_input = prompt_session.prompt(
                [("class:prompt", "  💬 你"), ("class:separator", " › ")],
            ).strip()
            if not user_input:
                continue

            result = session.process_message(user_input)
            if result == "__EXIT__":
                break

            console.print()
            console.print(Panel(
                Markdown(result),
                title="🤖 AI 分析师",
                border_style="cyan",
                padding=(1, 2),
            ))
            console.print()

        except KeyboardInterrupt:
            console.print("\n  [dim]再见！[/]")
            break
        except EOFError:
            break


# ================================================================
# 批量分析 + 导出
# ================================================================

def run_analysis_flow(data: list, data_file: Path, output_dir: Path = None):
    """分析流程：选模型 → 并行分析 → 报告 → AI问答"""
    from concurrent.futures import ThreadPoolExecutor, as_completed

    selected_models = select_analysis_models()
    if not selected_models:
        console.print("\n[yellow]未选择分析模型[/]")
        return

    console.print()
    console.print(Panel.fit(
        "[bold magenta]运行分析 & 生成报告[/]",
        border_style="magenta"
    ))

    quant_models = [m for m in selected_models if m in ANALYSIS_MODELS.get("quantitative", {})]
    qual_models = [m for m in selected_models if m in ANALYSIS_MODELS.get("qualitative", {})]

    results = {}
    with ThreadPoolExecutor(max_workers=6) as pool:
        futures = {}
        for model in quant_models:
            futures[pool.submit(_run_single_quant, data, model, data_file)] = model
        if qual_models and LLM_API_KEY:
            for model in qual_models:
                futures[pool.submit(_run_single_qual, data, model, data_file)] = model

        total = len(futures)
        done_count = 0
        for future in as_completed(futures):
            done_count += 1
            model = futures[future]
            try:
                result = future.result()
                results.update(result)
                status = "✓" if "error" not in list(result.values())[0] else "⚠"
                model_name = list(result.values())[0].get("model", model)
                console.print(f"  [{done_count}/{total}] {status} {model_name}")
            except Exception as e:
                console.print(f"  [{done_count}/{total}] ✗ {model}: {e}")

    if not results:
        console.print("\n[red]✗ 所有分析模型均失败[/]")
        return

    crawler_info = {"platform": "本地文件", "name": data_file.name}
    if output_dir is None:
        output_dir = data_file.parent

    console.print("\n  [cyan]📝 正在生成报告...[/]")
    report_path = generate_report(crawler_info, results, data_file, output_dir)
    console.print(f"  [green]✓[/] 报告已生成: [bold]{report_path}[/]")

    report_data = {
        "crawler": crawler_info,
        "analysis": results,
        "data_file": str(data_file),
        "generated_at": datetime.now().isoformat(),
    }

    json_path = report_path.with_suffix(".qa.json")
    json_path.write_text(json.dumps(report_data, ensure_ascii=False, indent=2, default=str), encoding="utf-8")

    if confirm_step("导出为 PDF？"):
        _export_pdf(report_path)

    if confirm_step("进入AI问答模式？"):
        chat_loop(report_path, report_data)

    _show_summary(report_path, data_file)
    return report_data


def _show_summary(report_path: Path, data_file: Path):
    console.clear()
    console.print()
    console.print(Panel.fit(
        "\n".join([
            "[bold green]✅ 分析完成！[/]",
            "",
            f"📊 报告文件: [bold cyan]{report_path}[/]",
            f"📁 数据文件: [bold cyan]{data_file}[/]",
        ]),
        border_style="green",
        title="[bold]任务总结[/]",
    ))
    console.print()
    if confirm_step("在浏览器中打开报告？"):
        import webbrowser
        webbrowser.open(f"file://{report_path}")


def _export_pdf(html_path: Path):
    """将 HTML 报告导出为 PDF"""
    from report.generator import export_pdf
    pdf_path = export_pdf(html_path)
    if pdf_path:
        console.print(f"  [green]✓[/] PDF 已生成: [bold]{pdf_path}[/]")
    else:
        console.print("  [yellow]⚠ PDF 导出失败，请安装 weasyprint: pip install weasyprint[/]")


# ================================================================
# 主流程
# ================================================================

def main():
    show_welcome()

    # 检测已保存的会话
    sessions = AnalystSession.list_sessions()
    if sessions:
        console.print()
        console.print(Panel.fit(
            "[bold yellow]📋 发现已保存的分析会话[/]",
            border_style="yellow"
        ))
        from rich.table import Table
        t = Table(show_header=True, header_style="bold")
        t.add_column("#", width=4)
        t.add_column("时间", width=18)
        t.add_column("数据文件", width=30)
        t.add_column("对话", width=8)
        for i, s in enumerate(sessions, 1):
            fname = Path(s["file"]).name if s["file"] else "?"
            t.add_row(str(i), s["saved"], fname, f"{s['messages']}轮")
        t.add_row("n", "—", "开始新分析", "—")
        console.print(t)
        console.print()

        choice = Prompt.ask("[bold]选择会话[/]", default="n")
        if choice.isdigit():
            idx = int(choice) - 1
            if 0 <= idx < len(sessions):
                _resume_session(sessions[idx]["id"])
                return

    mode = select_mode()
    if not mode:
        console.print("\n[yellow]已取消[/]")
        return

    save_mode_state(mode)

    # 延迟导入 cli/commands（避免循环引用）
    from cli.commands import (
        natural_language_crawl, run_skill_template, manage_schedule,
        start_web_ui, login_setup, check_and_login_if_needed,
        run_batch_analysis,
    )

    if mode == "nl_crawl":
        natural_language_crawl(load_data, run_ai_analysis_flow)
        return

    if mode == "template":
        run_skill_template(load_data, preview_data, run_ai_analysis_flow)
        return

    if mode == "schedule":
        manage_schedule()
        return

    if mode == "web":
        start_web_ui()
        return

    if mode == "login":
        login_setup()
        return

    if mode == "existing":
        if confirm_step("批量对比多个数据文件？"):
            files = select_multiple_files()
            if files:
                run_batch_analysis(files, load_data)
            return

        data_file = select_existing_file()
        if not data_file:
            console.print("\n[yellow]已取消[/]")
            return

        console.print(f"\n  [dim]正在加载数据...[/]")
        data = load_data(data_file)
        if not data:
            console.print("[red]✗ 数据为空或无法读取[/]")
            return

        if not preview_data(data, data_file):
            console.print("\n[yellow]已取消[/]")
            return

        run_ai_analysis_flow(data, data_file)
        return

    else:
        crawler_selection = select_crawler()
        if not crawler_selection:
            console.print("\n[yellow]已取消[/]")
            return

        crawler_key, crawler_info = crawler_selection
        save_crawler_state(crawler_key)

        args = get_crawler_args(crawler_info)

        check_and_login_if_needed(crawler_key)

        if not confirm_step("确认运行爬虫？"):
            console.print("\n[yellow]已取消[/]")
            return

        output_dir = run_crawler(crawler_key, args)
        files = get_latest_output_files(output_dir)

        if not files:
            console.print("\n[red]✗ 未找到输出文件[/]")
            return

        data_file = files[0]
        if len(files) > 1:
            console.print("\n  [bold]找到多个输出文件:[/]")
            for i, f in enumerate(files, 1):
                console.print(f"    {i}. {f.name}")
            try:
                choice = int(input("  选择: ")) - 1
                if 0 <= choice < len(files):
                    data_file = files[choice]
            except (ValueError, IndexError):
                pass

        console.print(f"\n  [green]✓[/] 数据文件: [bold]{data_file.name}[/]")

        console.print("\n  [dim]正在加载数据...[/]")
        data = load_data(data_file)
        if not data:
            console.print("[red]✗ 数据为空或无法读取[/]")
            return

        if not preview_data(data, data_file):
            console.print("\n[yellow]已取消[/]")
            return

        run_ai_analysis_flow(data, data_file)


def _resume_session(session_id: str):
    """恢复已保存的会话"""
    result = AnalystSession.load_session(session_id)
    if not result:
        console.print("[red]会话数据文件已不存在[/]")
        return
    data, data_file = result

    data = load_data(data_file)
    if not data:
        console.print("[red]无法加载数据文件[/]")
        return

    console.print(f"\n  [green]✓[/] 恢复会话: {data_file.name}")
    run_ai_analysis_flow(data, data_file)


def chat_loop(report_path: Path, report_data: dict):
    """基于报告的 AI 问答（旧模式，保留兼容）"""
    session = AnalystSession(
        [{"k": str(v)[:200] for k, v in item.items()} for item in [report_data]],
        report_path,
    )
    chat_loop_analyst(session)


if __name__ == "__main__":
    try:
        from config import OUTPUT_DIR, DATA_DIR, REPORT_DIR
        for dir_path in [OUTPUT_DIR, DATA_DIR, REPORT_DIR]:
            dir_path.mkdir(parents=True, exist_ok=True)
    except Exception:
        pass

    try:
        from validator import ConfigValidator
        if not ConfigValidator.print_status():
            console.print("\n[yellow]提示: 部分配置缺失，某些功能可能不可用[/]")

        main()
    except KeyboardInterrupt:
        console.print("\n\n[yellow]用户中断[/]")
        sys.exit(0)
    except Exception as e:
        logger.error(f"程序异常退出: {e}", exc_info=True)
        console.print(f"\n[red]发生错误: {e}[/]")
        sys.exit(1)