"""可视化输出模块

提供终端可视化功能，类似云端 SSE 实时流的本地替代。
"""

from __future__ import annotations

import shutil
import sys
from typing import Any, Callable, Optional

from .eval_types import EvalEvent, EvalEventType

class TerminalVisualizer:
    """终端可视化器

    将评测事件实时输出到终端，支持：
    - OA 序列可视化
    - 沙箱日志输出
    - 进度条显示
    - 彩色输出

    示例:
        visualizer = TerminalVisualizer(
            show_oa_sequence=True,
            show_user_logs=True,
        )

        for event in evaluator.evaluate_stream():
            visualizer.on_event(event)
    """

    # ANSI 颜色代码
    COLORS = {
        "reset": "\033[0m",
        "bold": "\033[1m",
        "dim": "\033[2m",
        "green": "\033[32m",
        "blue": "\033[34m",
        "yellow": "\033[33m",
        "red": "\033[31m",
        "cyan": "\033[36m",
        "magenta": "\033[35m",
        "gray": "\033[90m",
    }

    def __init__(
        self,
        show_oa_sequence: bool = True,
        show_judge_logs: bool = False,
        show_user_logs: bool = True,
        show_progress: bool = True,
        show_case_status: bool = True,
        colorize: bool = True,
        max_content_length: int = 200,
        output: Optional[Callable[[str], None]] = None,
    ):
        """初始化终端可视化器

        Args:
            show_oa_sequence: 是否显示 Observation-Action 序列
            show_judge_logs: 是否显示 Judge 沙箱输出
            show_user_logs: 是否显示 User 沙箱输出
            show_progress: 是否显示进度条
            show_case_status: 是否显示 Case 开始/结束状态
            colorize: 是否使用彩色输出
            max_content_length: 内容最大显示长度
            output: 自定义输出函数（默认 print）
        """
        self.show_oa_sequence = show_oa_sequence
        self.show_judge_logs = show_judge_logs
        self.show_user_logs = show_user_logs
        self.show_progress = show_progress
        self.show_case_status = show_case_status
        self.colorize = colorize and sys.stdout.isatty()
        self.max_content_length = max_content_length
        self.output = output or print

        # 状态追踪
        self._total_cases: Optional[int] = None
        self._completed_cases: set[int] = set()
        self._current_case: Optional[int] = None

    def _color(self, text: str, color: str) -> str:
        """添加颜色"""
        if not self.colorize:
            return text
        return f"{self.COLORS.get(color, '')}{text}{self.COLORS['reset']}"

    def _truncate(self, text: str) -> str:
        """截断长文本"""
        if len(text) <= self.max_content_length:
            return text
        return text[: self.max_content_length - 3] + "..."

    def _format_json_preview(self, data: dict) -> str:
        """格式化 JSON 预览"""
        import json

        try:
            text = json.dumps(data, ensure_ascii=False, default=str)
            return self._truncate(text)
        except Exception:
            return self._truncate(str(data))

    def on_event(self, event: EvalEvent) -> None:
        """处理评测事件

        Args:
            event: 评测事件
        """
        match event.type:
            case EvalEventType.CASE_START:
                self._on_case_start(event)
            case EvalEventType.CASE_END:
                self._on_case_end(event)
            case EvalEventType.OBSERVATION:
                self._on_observation(event)
            case EvalEventType.ACTION:
                self._on_action(event)
            case EvalEventType.JUDGE_LOG:
                self._on_judge_log(event)
            case EvalEventType.USER_LOG:
                self._on_user_log(event)
            case EvalEventType.PROGRESS:
                self._on_progress(event)
            case EvalEventType.ERROR:
                self._on_error(event)

    def _on_case_start(self, event: EvalEvent) -> None:
        """处理 Case 开始"""
        if not self.show_case_status:
            return

        self._current_case = event.case_index
        case_num = event.case_index + 1
        msg = self._color(f"┌── Case {case_num} ", "bold")
        self.output(msg)

    def _on_case_end(self, event: EvalEvent) -> None:
        """处理 Case 结束"""
        self._completed_cases.add(event.case_index)

        if not self.show_case_status:
            return

        status = event.data.get("status", "unknown")
        score = event.data.get("score", 0)
        error = event.data.get("error")

        # 根据状态选择颜色
        if status == "passed":
            status_color = "green"
            icon = "✓"
        elif status == "failed":
            status_color = "yellow"
            icon = "✗"
        else:
            status_color = "red"
            icon = "✗"

        case_num = event.case_index + 1
        status_text = self._color(f"{icon} Case {case_num}: {status.upper()}", status_color)

        if error:
            self.output(f"│  {status_text} (score: {score}) - Error: {error}")
        else:
            self.output(f"│  {status_text} (score: {score})")
        self.output(self._color("└" + "─" * 40, "dim"))

    def _on_observation(self, event: EvalEvent) -> None:
        """处理 Observation（Judge -> User）"""
        if not self.show_oa_sequence:
            return

        content = event.data.get("data", "")
        if isinstance(content, dict):
            preview = self._format_json_preview(content)
        else:
            preview = self._truncate(str(content))

        arrow = self._color("───>", "green")
        prefix = self._color("[OBS]", "green")
        self.output(f"│  {prefix} {arrow} {preview}")

    def _on_action(self, event: EvalEvent) -> None:
        """处理 Action（User -> Judge）"""
        if not self.show_oa_sequence:
            return

        content = event.data.get("data", "")
        if isinstance(content, dict):
            preview = self._format_json_preview(content)
        else:
            preview = self._truncate(str(content))

        arrow = self._color("<───", "blue")
        prefix = self._color("[ACT]", "blue")
        self.output(f"│  {prefix} {arrow} {preview}")

    def _on_judge_log(self, event: EvalEvent) -> None:
        """处理 Judge 日志"""
        if not self.show_judge_logs:
            return

        content = event.data.get("data", "")
        for line in str(content).strip().split("\n"):
            if line.strip():
                prefix = self._color("[JUDGE]", "magenta")
                self.output(f"│  {prefix} {self._truncate(line)}")

    def _on_user_log(self, event: EvalEvent) -> None:
        """处理 User 日志"""
        if not self.show_user_logs:
            return

        content = event.data.get("data", "")
        for line in str(content).strip().split("\n"):
            if line.strip():
                prefix = self._color("[USER]", "cyan")
                self.output(f"│  {prefix} {self._truncate(line)}")

    def _on_progress(self, event: EvalEvent) -> None:
        """处理进度更新"""
        completed = event.data.get("completed", 0)
        total = event.data.get("total", 0)
        self._total_cases = total

        if not self.show_progress or total <= 0:
            return

        # 绘制进度条
        terminal_width = shutil.get_terminal_size().columns
        bar_width = min(40, terminal_width - 30)

        filled = int(bar_width * completed / total)
        bar = "█" * filled + "░" * (bar_width - filled)

        percent = completed / total * 100
        progress_text = f"[{bar}] {completed}/{total} ({percent:.1f}%)"

        self.output(self._color(f"\n📊 Progress: {progress_text}\n", "bold"))

    def _on_error(self, event: EvalEvent) -> None:
        """处理错误"""
        error = event.data.get("error", "Unknown error")
        self.output(self._color(f"\n❌ Error: {error}\n", "red"))

    def print_summary(self, result: Any) -> None:
        """打印评测摘要

        Args:
            result: PhaseResult 对象
        """
        self.output("")
        self.output(self._color("=" * 50, "bold"))
        self.output(self._color("📋 Evaluation Summary", "bold"))
        self.output(self._color("=" * 50, "bold"))

        status = getattr(result, "status", "unknown")
        total_cases = getattr(result, "total_cases", 0)
        passed_cases = getattr(result, "passed_cases", 0)
        score = getattr(result, "score", 0)
        total_time = getattr(result, "total_time", 0)

        # 状态颜色
        if status == "success":
            status_str = self._color("✓ PASSED", "green")
        elif status == "failed":
            status_str = self._color("✗ FAILED", "yellow")
        else:
            status_str = self._color(f"✗ {status.upper()}", "red")

        self.output(f"Status:      {status_str}")
        self.output(f"Cases:       {passed_cases}/{total_cases} passed")
        self.output(f"Score:       {score}")
        self.output(f"Total Time:  {total_time}ms")

        # 每个 case 的结果
        cases = getattr(result, "cases", [])
        if cases:
            self.output("")
            self.output(self._color("Case Details:", "bold"))
            for case in cases:
                case_status = getattr(case, "status", "unknown")
                case_score = getattr(case, "score", 0)
                case_idx = getattr(case, "case_index", 0) + 1

                if case_status == "passed":
                    case_icon = self._color("✓", "green")
                elif case_status == "failed":
                    case_icon = self._color("✗", "yellow")
                else:
                    case_icon = self._color("✗", "red")

                self.output(f"  {case_icon} Case {case_idx}: {case_status} (score: {case_score})")

        self.output(self._color("=" * 50, "bold"))
