from __future__ import annotations

import datetime
import time

from pydantic_ai import Agent
from pydantic_ai.models import Model

from core.models import DiscoveryResult
from generator.strategies.mutation import MutationStrategy
from generator.strategies.objective import ObjectiveStrategy
from goals.base import GoalResult
from goals.objective import ObjectiveGoal
from output.base import AbstractOutput
from prompts.loader import load_probes
from prompts.models import Probe

_TOP_K = 5
_MIN_SCORE_FOR_SEED = 0.1
_PROMPT_PREVIEW = 600
_RESPONSE_PREVIEW = 800
_LINE = "━" * 62


def _log(outputs: list[AbstractOutput], msg: str) -> None:
    for o in outputs:
        o.log(msg)


def _status(outputs: list[AbstractOutput], msg: str) -> None:
    for o in outputs:
        o.update_status(msg)


def _score_bar(score: float, width: int = 12) -> str:
    filled = round(score * width)
    empty = width - filled
    if score >= 0.7:
        color = "green"
    elif score >= 0.35:
        color = "yellow"
    else:
        color = "red"
    return f"[{color}]{'█' * filled}[/{color}][dim]{'░' * empty}[/dim]"


def _model_name(model: Model) -> str:
    try:
        return model.model_name
    except Exception:
        return type(model).__name__


def _elapsed(t0: float) -> str:
    return f"{time.monotonic() - t0:.1f}s"


def _preview(text: str, limit: int) -> str:
    text = text.strip()
    return text[:limit] + ("…" if len(text) > limit else "")


def _block(text: str, indent: str = "    ", width: int = 88) -> list[str]:
    """Word-wrap text into indented lines for the log."""
    lines: list[str] = []
    for raw_line in text.splitlines():
        raw_line = raw_line.strip()
        if not raw_line:
            lines.append("")
            continue
        words = raw_line.split()
        current = indent
        for word in words:
            if len(current) + len(word) + 1 > width and current.strip():
                lines.append(current)
                current = indent + word
            else:
                current = (current + " " + word) if current.strip() else indent + word
        if current.strip():
            lines.append(current)
    return lines


class DiscoveryLoop:
    def __init__(
        self,
        generator_model: Model,
        target_model: Model,
        judge_model: Model,
        objective: str,
        outputs: list[AbstractOutput],
        iterations: int = 5,
        candidates_per_iter: int = 5,
    ) -> None:
        self._generator_model = generator_model
        self._target_model = target_model
        self._judge_model = judge_model
        self._objective = objective
        self._outputs = outputs
        self._iterations = iterations
        self._candidates_per_iter = candidates_per_iter

    def run(self) -> list[DiscoveryResult]:
        seeds = load_probes()
        active_seeds = list(seeds)
        all_results: list[DiscoveryResult] = []
        target_agent = Agent(self._target_model)
        goal = ObjectiveGoal(self._objective, self._judge_model)
        objective_strategy = ObjectiveStrategy([self._objective])
        mutation_strategy = MutationStrategy()
        total = self._iterations * self._candidates_per_iter
        probe_index = 0

        gen_name = _model_name(self._generator_model)
        tgt_name = _model_name(self._target_model)
        jdg_name = _model_name(self._judge_model)

        # ── run header ────────────────────────────────────────────────────────
        _log(self._outputs, "")
        _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
        _log(self._outputs, "[bold]  probegpt  —  automated red-team discovery[/bold]")
        _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
        _log(self._outputs, "")
        _log(self._outputs, f"  [bold]objective[/bold]   {self._objective}")
        _log(self._outputs, f"  [bold]target[/bold]      {tgt_name}")
        _log(self._outputs, f"  [bold]generator[/bold]   {gen_name}")
        _log(self._outputs, f"  [bold]judge[/bold]       {jdg_name}")
        _log(
            self._outputs,
            f"  [bold]plan[/bold]        {self._iterations} iters × {self._candidates_per_iter} candidates = {total} probes",
        )
        _log(self._outputs, f"  [bold]seeds[/bold]       {len(seeds)} probes loaded")
        _log(self._outputs, "")

        run_t0 = time.monotonic()

        for iteration in range(1, self._iterations + 1):
            iter_results: list[DiscoveryResult] = []

            # ── iteration header ───────────────────────────────────────────────
            _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
            _log(self._outputs, f"[bold]  ITERATION {iteration} / {self._iterations}[/bold]")
            _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
            _log(self._outputs, "")
            _log(self._outputs, f"  [dim]active seeds ({len(active_seeds)}):[/dim]")
            for s in active_seeds[:6]:
                _log(self._outputs, f"    [dim]· {s.title}[/dim]")
            if len(active_seeds) > 6:
                _log(self._outputs, f"    [dim]  … {len(active_seeds) - 6} more[/dim]")
            _log(self._outputs, "")

            # ── generation ────────────────────────────────────────────────────
            if iteration == 1:
                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · generating {self._candidates_per_iter} probes…",
                )
                _log(
                    self._outputs, f"  [bold]GENERATE[/bold]  strategy=objective  model={gen_name}"
                )
                _log(
                    self._outputs,
                    f"  [dim]writing {self._candidates_per_iter} probes for objective…[/dim]",
                )
                gen_t0 = time.monotonic()
                candidates = objective_strategy.generate(
                    active_seeds, self._generator_model, count=self._candidates_per_iter
                )
                strategy_names = ["objective"] * len(candidates)
                _log(
                    self._outputs,
                    f"  [dim]  ↳ {len(candidates)} probes ready  ({_elapsed(gen_t0)})[/dim]",
                )

            else:
                half = self._candidates_per_iter // 2
                rest = self._candidates_per_iter - half

                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · generating mutation probes…",
                )
                _log(
                    self._outputs,
                    f"  [bold]GENERATE[/bold]  strategy=mutation+objective  model={gen_name}",
                )
                _log(
                    self._outputs, f"  [dim]mutation: writing {half} variants from top seeds…[/dim]"
                )
                gen_t0 = time.monotonic()
                mut_candidates = mutation_strategy.generate(
                    active_seeds, self._generator_model, count=half
                )
                _log(
                    self._outputs,
                    f"  [dim]  ↳ {len(mut_candidates)} mutation probes  ({_elapsed(gen_t0)})[/dim]",
                )

                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · generating objective probes…",
                )
                _log(self._outputs, f"  [dim]objective: writing {rest} probes for goal…[/dim]")
                obj_t0 = time.monotonic()
                obj_candidates = objective_strategy.generate(
                    active_seeds, self._generator_model, count=rest
                )
                _log(
                    self._outputs,
                    f"  [dim]  ↳ {len(obj_candidates)} objective probes  ({_elapsed(obj_t0)})[/dim]",
                )

                candidates = mut_candidates + obj_candidates
                strategy_names = ["mutation"] * len(mut_candidates) + ["objective"] * len(
                    obj_candidates
                )

            _log(self._outputs, "")

            # ── probe loop ────────────────────────────────────────────────────
            for prompt, strategy_name in zip(candidates, strategy_names):
                probe_index += 1

                # probe header
                _log(
                    self._outputs,
                    f"[dim]  ┌─ probe {probe_index}/{total}  [{strategy_name}] {'─' * max(0, 44 - len(str(probe_index)) - len(str(total)) - len(strategy_name))}[/dim]",
                )
                _log(self._outputs, "")

                # full prompt
                _log(self._outputs, "  [bold yellow]PROMPT[/bold yellow]")
                _log(self._outputs, "  [dim]┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄[/dim]")
                for line in _block(_preview(prompt, _PROMPT_PREVIEW)):
                    _log(self._outputs, line)
                _log(self._outputs, "")

                # target call
                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · probe {probe_index}/{total} · target…",
                )
                _log(self._outputs, f"  [dim]→ {tgt_name}…[/dim]")
                tgt_t0 = time.monotonic()
                try:
                    target_result = target_agent.run_sync(prompt)
                    response_text: str = target_result.output
                except Exception as exc:
                    response_text = f"[ERROR: {exc}]"
                tgt_elapsed = _elapsed(tgt_t0)

                _log(self._outputs, f"  [bold cyan]RESPONSE[/bold cyan]  [dim]{tgt_elapsed}[/dim]")
                _log(self._outputs, "  [dim]┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄[/dim]")
                for line in _block(_preview(response_text, _RESPONSE_PREVIEW)):
                    _log(self._outputs, line)
                _log(self._outputs, "")

                # judge call
                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · probe {probe_index}/{total} · judging…",
                )
                _log(self._outputs, f"  [dim]⚖ {jdg_name}…[/dim]")
                jdg_t0 = time.monotonic()
                try:
                    goal_result = goal.evaluate(prompt=prompt, response=response_text)
                except Exception as exc:
                    goal_result = GoalResult(
                        achieved=False,
                        score=0.0,
                        reason=f"[EVAL ERROR: {exc}]",
                        goal_name="objective",
                    )
                jdg_elapsed = _elapsed(jdg_t0)

                bar = _score_bar(goal_result.score)
                if goal_result.achieved:
                    verdict_tag = "[bold green]  ✓ ACHIEVED  [/bold green]"
                else:
                    verdict_tag = "[dim]✗ not achieved[/dim]"

                _log(
                    self._outputs,
                    f"  [bold magenta]VERDICT[/bold magenta]  [dim]{jdg_elapsed}[/dim]",
                )
                _log(self._outputs, "  [dim]┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄[/dim]")
                _log(self._outputs, f"  {verdict_tag}  {bar}  {goal_result.score:.2f}")
                _log(self._outputs, f"  [italic dim]{goal_result.reason}[/italic dim]")
                _log(self._outputs, "")
                _log(self._outputs, f"[dim]  └{'─' * 55}[/dim]")
                _log(self._outputs, "")

                r = DiscoveryResult(
                    iteration=iteration,
                    strategy_name=strategy_name,
                    goal_name="objective",
                    prompt=prompt,
                    response_content=response_text,
                    succeeded=goal_result.achieved,
                    score=goal_result.score,
                    seed_titles=[s.title for s in active_seeds[:3]],
                    timestamp=datetime.datetime.utcnow().isoformat(),
                )
                for output in self._outputs:
                    output.emit(r)
                iter_results.append(r)

            all_results.extend(iter_results)

            # ── iteration summary ─────────────────────────────────────────────
            n_succeeded = sum(1 for r in iter_results if r.succeeded)
            best = max((r.score for r in iter_results), default=0.0)
            avg = sum(r.score for r in iter_results) / len(iter_results) if iter_results else 0.0

            top_k = sorted(
                [r for r in iter_results if r.score > _MIN_SCORE_FOR_SEED],
                key=lambda r: r.score,
                reverse=True,
            )[:_TOP_K]

            _log(self._outputs, f"  [bold]ITER {iteration} SUMMARY[/bold]")
            _log(self._outputs, f"    probes     {len(iter_results)}")
            _log(self._outputs, f"    achieved   {n_succeeded} / {len(iter_results)}")
            _log(self._outputs, f"    best       {_score_bar(best)} {best:.2f}")
            _log(self._outputs, f"    avg        {_score_bar(avg)} {avg:.2f}")

            if top_k and iteration < self._iterations:
                _log(self._outputs, f"    top-{len(top_k)} → next iter:")
                for r in top_k:
                    _log(
                        self._outputs,
                        f"      [dim]{_score_bar(r.score, 6)} {r.score:.2f}  {_preview(r.prompt, 70)}[/dim]",
                    )

            _log(self._outputs, "")

            if top_k:
                new_seeds = [
                    Probe(
                        title=f"[iter-{r.iteration}] {r.strategy_name}",
                        content=r.prompt,
                        category="generated",
                    )
                    for r in top_k
                ]
                active_seeds = (active_seeds + new_seeds)[-20:]

        # ── final summary ─────────────────────────────────────────────────────
        total_elapsed = _elapsed(run_t0)
        n_total = len(all_results)
        n_achieved = sum(1 for r in all_results if r.succeeded)
        overall_best = max((r.score for r in all_results), default=0.0)

        _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
        _log(self._outputs, "[bold]  FINAL RESULTS[/bold]")
        _log(self._outputs, f"[bold cyan]{_LINE}[/bold cyan]")
        _log(self._outputs, "")
        _log(self._outputs, f"  probes      {n_total}")
        _log(
            self._outputs,
            f"  achieved    {n_achieved} / {n_total}  ({n_achieved / n_total * 100:.0f}%)"
            if n_total
            else "  achieved    0",
        )
        _log(self._outputs, f"  best score  {_score_bar(overall_best)} {overall_best:.2f}")
        _log(self._outputs, f"  elapsed     {total_elapsed}")
        _log(self._outputs, "")

        if n_achieved:
            _log(self._outputs, "  [bold green]top probes:[/bold green]")
            top_all = sorted(
                [r for r in all_results if r.succeeded],
                key=lambda r: r.score,
                reverse=True,
            )[:5]
            for i, r in enumerate(top_all, 1):
                _log(
                    self._outputs,
                    f"  [bold]{i}.[/bold]  score={r.score:.2f}  iter={r.iteration}  [{r.strategy_name}]",
                )
                _log(self._outputs, f"     [dim]{_preview(r.prompt, 120)}[/dim]")
                _log(self._outputs, "")

        for output in self._outputs:
            output.finalize(all_results)

        return all_results
