#!/usr/bin/env python3
"""控制台问数脚本：支持交互式多轮问答和单次执行。"""

from __future__ import annotations

import argparse
import json
import sys
import traceback
import uuid
from pathlib import Path
from typing import Any, Dict, Iterable, Optional, Sequence

# 允许直接使用 `python scripts/console_chat.py` 运行
PROJECT_ROOT = Path(__file__).resolve().parent.parent
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from agent.agents.conversation import ConversationAgent
from config import Config


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        description="在控制台中直接输入自然语言问题，调用 ConversationAgent 进行问数。"
    )
    parser.add_argument("--model-section", default=None, help="可选：指定 config.ini 中的模型配置段")
    parser.add_argument("--conversation-id", default=None, help="可选：会话 ID，仅用于展示和上下文标识")
    parser.add_argument("--user", default="console-user", help="可选：用户标识")
    parser.add_argument("--query", "-q", default=None, help="单次执行模式：直接执行一条问题后退出")
    parser.add_argument("--skip-sr-api", action="store_true", help="仅生成 SQL，不执行 SR API")
    parser.add_argument("--show-sql", action="store_true", help="额外打印生成的 SQL")
    parser.add_argument("--show-context", action="store_true", help="额外打印完整上下文 JSON")
    parser.add_argument("--show-plan", action="store_true", help="额外打印 SQL 规划 JSON")
    parser.add_argument("--no-banner", action="store_true", help="不显示启动横幅")
    return parser


def print_banner(model_section: Optional[str], conversation_id: str, skip_sr_api: bool) -> None:
    print("=" * 72)
    print("APBO Console Chat")
    print(f"Model Section : {model_section or 'default'}")
    print(f"Conversation  : {conversation_id}")
    print(f"Execution     : {'SQL only' if skip_sr_api else 'SQL + SR API'}")
    print("Commands      : /quit /exit /sql /context /plan /exec")
    print("=" * 72)


def _safe_json(data: Any) -> str:
    try:
        return json.dumps(data, ensure_ascii=False, indent=2, default=str)
    except Exception:
        return str(data)


def _looks_like_json(text: str) -> bool:
    stripped = (text or "").strip()
    return stripped.startswith("{") or stripped.startswith("[")


def _try_json_loads(value: Any) -> Any:
    if isinstance(value, (dict, list)):
        return value
    if isinstance(value, str) and _looks_like_json(value):
        try:
            return json.loads(value)
        except Exception:
            return value
    return value


def parse_sr_api_result(raw_result: Any) -> Any:
    """解析 SR API 返回值，兼容外层 envelope 和内层 text JSON。"""
    parsed = _try_json_loads(raw_result)
    if isinstance(parsed, dict) and "text" in parsed:
        text_payload = _try_json_loads(parsed.get("text"))
        parsed = {**parsed, "text": text_payload}
    return parsed


def _find_table_candidate(value: Any) -> Any:
    if isinstance(value, list):
        return value
    if isinstance(value, dict):
        if "columns" in value and any(key in value for key in ("rows", "data", "values")):
            return value
        for key in ("data", "rows", "records", "record", "items", "list", "result", "text"):
            nested = value.get(key)
            if isinstance(nested, (list, dict)):
                found = _find_table_candidate(nested)
                if found is not None:
                    return found
    return None


def extract_table_rows(parsed_result: Any) -> tuple[list[str], list[list[str]]]:
    """从常见查询结果结构中提取表头和二维行数据。"""
    candidate = _find_table_candidate(parsed_result)
    if candidate is None:
        return [], []

    if isinstance(candidate, dict) and isinstance(candidate.get("columns"), list):
        headers = [str(col) for col in candidate.get("columns") or []]
        raw_rows = candidate.get("rows") or candidate.get("data") or candidate.get("values") or []
        if raw_rows and all(isinstance(row, dict) for row in raw_rows):
            return headers or list(raw_rows[0].keys()), [
                [str((row or {}).get(header, "")) for header in (headers or list(raw_rows[0].keys()))]
                for row in raw_rows
            ]
        return headers, [[str(cell) for cell in row] for row in raw_rows if isinstance(row, (list, tuple))]

    if isinstance(candidate, list) and candidate:
        if all(isinstance(row, dict) for row in candidate):
            headers: list[str] = []
            for row in candidate:
                for key in row.keys():
                    if key not in headers:
                        headers.append(str(key))
            return headers, [[str((row or {}).get(header, "")) for header in headers] for row in candidate]
        if all(isinstance(row, (list, tuple)) for row in candidate):
            width = max(len(row) for row in candidate)
            headers = [f"col_{idx + 1}" for idx in range(width)]
            return headers, [[str(row[idx]) if idx < len(row) else "" for idx in range(width)] for row in candidate]

    return [], []


def render_text_table(headers: Sequence[str], rows: Sequence[Sequence[str]], *, max_width: int = 28, max_rows: int = 20) -> str:
    """将二维数据渲染成纯文本表格。"""
    if not headers or not rows:
        return "<empty table>"

    def clip(value: Any) -> str:
        text = str(value).replace("\r", " ").replace("\n", " ")
        return text if len(text) <= max_width else text[: max_width - 3] + "..."

    display_rows = list(rows[:max_rows])
    str_rows = [[clip(cell) for cell in row] for row in display_rows]
    clipped_headers = [clip(header) for header in headers]

    widths = []
    for idx, header in enumerate(clipped_headers):
        col_values = [row[idx] if idx < len(row) else "" for row in str_rows]
        widths.append(max(len(header), *(len(value) for value in col_values)) if col_values else len(header))

    def render_row(values: Sequence[str]) -> str:
        padded = []
        for idx, width in enumerate(widths):
            value = values[idx] if idx < len(values) else ""
            padded.append(value.ljust(width))
        return "| " + " | ".join(padded) + " |"

    separator = "+-" + "-+-".join("-" * width for width in widths) + "-+"
    lines = [separator, render_row(clipped_headers), separator]
    lines.extend(render_row(row) for row in str_rows)
    lines.append(separator)
    if len(rows) > max_rows:
        lines.append(f"... showing first {max_rows} of {len(rows)} rows")
    return "\n".join(lines)


def _format_sr_api_table(sr_api_result: Any) -> Optional[str]:
    parsed = parse_sr_api_result(sr_api_result)
    headers, rows = extract_table_rows(parsed)
    if not headers or not rows:
        return None

    summary = f"Query Result: {len(rows)} row(s)"
    if isinstance(parsed, dict) and parsed.get("status_code") is not None:
        summary += f" | status={parsed.get('status_code')}"
    return summary + "\n" + render_text_table(headers, rows)


def format_result(
    result: Dict[str, Any],
    *,
    show_sql: bool = False,
    show_context: bool = False,
    show_plan: bool = False,
) -> str:
    context = (result or {}).get("context") or {}
    messages = (result or {}).get("messages") or []
    sr_api_result = context.get("sr_api_result")

    answer = ""
    if messages:
        last = messages[-1]
        answer = getattr(last, "content", "") or str(last)
    if not answer:
        answer = str(sr_api_result or context.get("final_sql") or "<empty response>")

    table_block = _format_sr_api_table(sr_api_result) if sr_api_result else None
    if table_block and (_looks_like_json(answer) or answer.strip().lower() in {"ok", "success"}):
        answer = "查询成功，结果已按二维表格展示如下。"

    blocks = [f"Answer:\n{answer}"]

    if table_block:
        blocks.append(table_block)

    if show_sql and context.get("final_sql"):
        blocks.append(f"SQL:\n{context['final_sql']}")

    if show_plan and context.get("sql_plan"):
        blocks.append(f"SQL Plan:\n{_safe_json(context['sql_plan'])}")

    if show_context:
        blocks.append(f"Context:\n{_safe_json(context)}")

    return "\n\n".join(blocks)


def run_turn(
    agent: ConversationAgent,
    query: str,
    *,
    user: str,
    conversation_id: str,
    skip_sr_api: bool,
    show_sql: bool,
    show_context: bool,
    show_plan: bool,
) -> Dict[str, Any]:
    result = agent.run(
        query,
        user=user,
        conversation_id=conversation_id,
        skip_sr_api=skip_sr_api,
        debug_node_trace=True,
    )
    print(format_result(result, show_sql=show_sql, show_context=show_context, show_plan=show_plan))
    return result


def interactive_loop(args: argparse.Namespace) -> int:
    conversation_id = args.conversation_id or f"console_{uuid.uuid4().hex[:8]}"
    agent = ConversationAgent(model_section=args.model_section)

    show_sql = bool(args.show_sql)
    show_context = bool(args.show_context)
    show_plan = bool(args.show_plan)
    skip_sr_api = bool(args.skip_sr_api)

    if not args.no_banner:
        print_banner(args.model_section, conversation_id, skip_sr_api)

    while True:
        try:
            user_input = input("\n问数> ").strip()
        except (EOFError, KeyboardInterrupt):
            print("\nBye.")
            return 0

        if not user_input:
            continue

        lowered = user_input.lower()
        if lowered in {"/quit", "/exit", "quit", "exit"}:
            print("Bye.")
            return 0
        if lowered == "/sql":
            show_sql = not show_sql
            print(f"show_sql = {show_sql}")
            continue
        if lowered == "/context":
            show_context = not show_context
            print(f"show_context = {show_context}")
            continue
        if lowered == "/plan":
            show_plan = not show_plan
            print(f"show_plan = {show_plan}")
            continue
        if lowered == "/exec":
            skip_sr_api = not skip_sr_api
            print(f"skip_sr_api = {skip_sr_api}")
            continue

        try:
            run_turn(
                agent,
                user_input,
                user=args.user,
                conversation_id=conversation_id,
                skip_sr_api=skip_sr_api,
                show_sql=show_sql,
                show_context=show_context,
                show_plan=show_plan,
            )
        except Exception as exc:
            print(f"\n[ERROR] {exc}")
            traceback.print_exc()

    return 0


def one_shot(args: argparse.Namespace) -> int:
    conversation_id = args.conversation_id or f"console_{uuid.uuid4().hex[:8]}"
    agent = ConversationAgent(model_section=args.model_section)
    try:
        run_turn(
            agent,
            args.query,
            user=args.user,
            conversation_id=conversation_id,
            skip_sr_api=bool(args.skip_sr_api),
            show_sql=bool(args.show_sql),
            show_context=bool(args.show_context),
            show_plan=bool(args.show_plan),
        )
        return 0
    except Exception as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        traceback.print_exc()
        return 1


def main(argv: Optional[Iterable[str]] = None) -> int:
    parser = build_parser()
    args = parser.parse_args(list(argv) if argv is not None else None)

    try:
        Config.validate_config()
    except Exception as exc:
        print(f"Configuration error: {exc}", file=sys.stderr)
        print("Please check `config/config.ini` and your model/API settings.", file=sys.stderr)
        return 1

    if args.query:
        return one_shot(args)
    return interactive_loop(args)


if __name__ == "__main__":
    raise SystemExit(main())