from __future__ import annotations

from typing import Any

from rich.console import Console
from rich.progress import Progress, SpinnerColumn, TextColumn
from rich.table import Table

from leshy.core.models import (
    AttackPayload,
    EvalResult,
    TargetResponse,
    TestCaseConfig,
    TestCaseResult,
    TestSuiteConfig,
    TestSuiteResult,
    Verdict,
)
from leshy.core.registry import logger_registry
from leshy.loggers.base import Logger

_VERDICT_COLORS = {
    Verdict.VULNERABLE: "red bold",
    Verdict.RESILIENT: "green",
    Verdict.INCONCLUSIVE: "yellow",
    Verdict.ERROR: "magenta",
}


@logger_registry.register("console")
class ConsoleLogger(Logger):
    """Rich console output with color-coded verdicts and summary tables."""

    name = "console"

    def __init__(self, **kwargs: Any) -> None:
        super().__init__(**kwargs)
        self._console = Console()
        self._verbose = kwargs.get("verbose", False)
        self._progress: Progress | None = None
        self._task_id: Any = None
        self._payload_count: int = 0
        self._warned_parse: bool = False

    async def on_suite_start(self, config: TestSuiteConfig) -> None:
        self._console.print()
        self._console.rule(f"[bold blue]Leshy Security Scan: {config.name}")
        self._console.print(f"  Target: {config.target.url}")
        self._console.print(f"  Tests:  {len(config.tests)}")
        self._console.print()

    async def on_case_start(self, case: TestCaseConfig) -> None:
        mode = "multi-turn" if case.multi_turn else "single-turn"
        attacks = ", ".join(a.name for a in case.attacks)
        self._console.print(f"[bold]>> {case.name}[/bold] ({mode}) — attacks: {attacks}")

        self._warned_parse = False
        if not self._verbose:
            self._payload_count = 0
            self._progress = Progress(
                SpinnerColumn(),
                TextColumn("[progress.description]{task.description}"),
                console=self._console,
            )
            self._progress.start()
            self._task_id = self._progress.add_task("Processing...", total=None)

    async def on_payload_result(
        self,
        payload: AttackPayload,
        response: TargetResponse,
        results: list[EvalResult],
    ) -> None:
        if not self._warned_parse and response.metadata.get("parse_warning"):
            if self._progress is not None:
                self._progress.stop()
            self._console.print(
                f"  [yellow]Warning: {response.metadata['parse_warning']}[/yellow]"
            )
            if self._progress is not None:
                self._progress.start()
            self._warned_parse = True

        if self._verbose:
            for r in results:
                color = _VERDICT_COLORS.get(r.verdict, "white")
                truncated = payload.content[:80] + ("..." if len(payload.content) > 80 else "")
                self._console.print(
                    f"  [{color}]{r.verdict.value:>13}[/{color}] "
                    f"[dim]{payload.attack_name}[/dim] {truncated}"
                )
        elif self._progress is not None and self._task_id is not None:
            self._payload_count += 1
            self._progress.update(
                self._task_id,
                description=f"Processing... {self._payload_count} payloads",
            )

    async def on_case_complete(self, result: TestCaseResult) -> None:
        if self._progress is not None:
            self._progress.stop()
            self._progress = None
            self._task_id = None

        table = Table(show_header=False, box=None, padding=(0, 2))
        table.add_column("label", style="dim")
        table.add_column("value")
        table.add_row("Payloads", str(result.total_payloads))
        table.add_row("Vulnerable", f"[red bold]{result.vulnerable_count}[/red bold]")
        table.add_row("Resilient", f"[green]{result.resilient_count}[/green]")
        if result.inconclusive_count:
            table.add_row("Inconclusive", f"[yellow]{result.inconclusive_count}[/yellow]")
        if result.error_count:
            table.add_row("Errors", f"[magenta]{result.error_count}[/magenta]")
        latencies = [
            r.response_latency_ms for r in result.eval_results
            if r.response_latency_ms > 0
        ]
        if latencies:
            avg = sum(latencies) / len(latencies)
            table.add_row(
                "Avg response",
                f"[dim]{avg:.0f}ms[/dim]",
            )
        table.add_row("Duration", f"{result.duration_ms:.0f}ms")
        self._console.print(table)
        self._console.print()

    async def on_suite_complete(self, result: TestSuiteResult) -> None:
        self._console.rule("[bold blue]Summary")
        table = Table(title="Test Results")
        table.add_column("Test Case", style="bold")
        table.add_column("Payloads", justify="right")
        table.add_column("Vulnerable", justify="right", style="red bold")
        table.add_column("Resilient", justify="right", style="green")
        table.add_column("Errors", justify="right", style="magenta")
        table.add_column("Duration", justify="right", style="dim")

        for case in result.case_results:
            table.add_row(
                case.name,
                str(case.total_payloads),
                str(case.vulnerable_count),
                str(case.resilient_count),
                str(case.error_count),
                f"{case.duration_ms:.0f}ms",
            )

        self._console.print(table)
        self._console.print()

        if result.total_vulnerable > 0:
            self._console.print(
                f"[red bold]VULNERABLE:[/red bold] "
                f"{result.total_vulnerable}/{result.total_payloads} "
                f"payloads triggered vulnerabilities"
            )
        else:
            self._console.print("[green bold]CLEAN:[/green bold] No vulnerabilities detected")

        self._print_stats(result)
        self._console.print()

    def _print_stats(self, result: TestSuiteResult) -> None:
        """Print aggregate stats: latency, throughput, token usage."""
        latencies = [
            r.response_latency_ms
            for c in result.case_results
            for r in c.eval_results
            if r.response_latency_ms > 0
        ]
        if latencies:
            avg = sum(latencies) / len(latencies)
            lo, hi = min(latencies), max(latencies)
            total_ms = sum(c.duration_ms for c in result.case_results)
            pps = (
                f"{result.total_payloads / (total_ms / 1000):.1f}"
                if total_ms > 0
                else "—"
            )
            self._console.print(
                f"[dim]Response time: avg {avg:.0f}ms, "
                f"min {lo:.0f}ms, max {hi:.0f}ms  |  "
                f"{pps} payloads/sec[/dim]"
            )

        if result.token_usage.total_tokens > 0:
            self._console.print(
                f"[dim]Tokens used: {result.token_usage.total_tokens:,} "
                f"(prompt: {result.token_usage.prompt_tokens:,}, "
                f"completion: {result.token_usage.completion_tokens:,})"
                "[/dim]"
            )
