from __future__ import annotations

import datetime
import time
import traceback
from typing import TYPE_CHECKING

_FILTER_KEYWORDS = (
    "content_filter",
    "jailbreak",
    "content management policy",
    "ResponsibleAIPolicyViolation",
)


def _is_content_filtered(exc: Exception) -> bool:
    """Return True if the exception is an Azure/OpenAI content-filter block."""
    s = str(exc)
    return any(kw.lower() in s.lower() for kw in _FILTER_KEYWORDS)


from pydantic_ai import Agent
from pydantic_ai.models import Model

from probegpt.core.config import StrategyConfig
from probegpt.core.models import DiscoveryResult, Probe
from probegpt.data import load_seeds
from probegpt.grading import GradedSeed, SeedGrader, SeedPool
from probegpt.judge.models import GoalResult
from probegpt.judge.objective import ObjectiveJudge
from probegpt.output.base import AbstractOutput
from probegpt.strategies.distraction import DistractionStrategy
from probegpt.strategies.encoding import EncodingStrategy
from probegpt.strategies.linguistic import LinguisticStrategy
from probegpt.strategies.mutation import MutationStrategy
from probegpt.strategies.objective import ObjectiveStrategy
from probegpt.strategies.persuasion import PersuasionStrategy
from probegpt.strategies.roleplay import RoleplayStrategy
from probegpt.strategies.structural import StructuralStrategy

if TYPE_CHECKING:
    from probegpt.strategies.base import AbstractStrategy

_GRADING_SCORE_THRESHOLD = 0.5
_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
    color = "green" if score >= 0.7 else ("yellow" if score >= 0.35 else "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]:
    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


def _allocate(total: int, weights: dict[str, float]) -> dict[str, int]:
    """Distribute total candidates across strategy names proportionally to weights."""
    if not weights:
        return {}
    active = {k: v for k, v in weights.items() if v > 0}
    if not active:
        return {}

    total_w = sum(active.values())
    raw = {k: total * v / total_w for k, v in active.items()}
    floors = {k: int(v) for k, v in raw.items()}
    remainder = total - sum(floors.values())
    fractions = sorted(active.keys(), key=lambda k: -(raw[k] - floors[k]))
    for k in fractions[:remainder]:
        floors[k] += 1
    return {k: v for k, v in floors.items() if v > 0}


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,
        strategy_config: StrategyConfig | None = None,
    ) -> 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
        self._strategy_config = strategy_config or StrategyConfig()

    def _build_strategies(self) -> dict[str, AbstractStrategy]:
        obj = self._objective
        return {
            "objective": ObjectiveStrategy([obj]),
            "mutation": MutationStrategy(),
            "encoding": EncodingStrategy(obj),
            "roleplay": RoleplayStrategy(obj),
            "persuasion": PersuasionStrategy(obj),
            "linguistic": LinguisticStrategy(obj),
            "structural": StructuralStrategy(obj),
            "distraction": DistractionStrategy(obj),
        }

    def _active_weights(self) -> dict[str, float]:
        cfg = self._strategy_config
        return {
            name: cfg.effective_weight(name)
            for name in cfg.enabled
            if cfg.effective_weight(name) > 0
        }

    def run(self) -> list[DiscoveryResult]:
        seeds = load_seeds()
        pool = SeedPool(initial_seeds=seeds)
        grader = SeedGrader(model=self._judge_model, objective=self._objective)
        all_results: list[DiscoveryResult] = []
        target_agent = Agent(self._target_model)
        judge = ObjectiveJudge(self._objective, self._judge_model)
        strategies = self._build_strategies()
        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)
        enabled_names = self._strategy_config.enabled

        _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, f"  [bold]strategies[/bold]  {', '.join(enabled_names)}")
        _log(self._outputs, "")

        run_t0 = time.monotonic()

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

            _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, "")

            all_probes = pool.all_probes
            _log(self._outputs, f"  [dim]active seeds ({len(all_probes)}):[/dim]")
            for s in all_probes[:6]:
                _log(self._outputs, f"    [dim]· {s.title}[/dim]")
            if len(all_probes) > 6:
                _log(self._outputs, f"    [dim]  … {len(all_probes) - 6} more[/dim]")

            if technique_hints:
                tags_str = "  ".join(f"[cyan]{t}[/cyan]" for t in technique_hints[:5])
                _log(self._outputs, "")
                _log(self._outputs, f"  [dim]top techniques:  {tags_str}[/dim]")
            _log(self._outputs, "")

            active_weights = self._active_weights()
            allocation = _allocate(self._candidates_per_iter, active_weights)

            if not allocation:
                _log(
                    self._outputs, "  [yellow]  no active strategies — skipping iteration[/yellow]"
                )
                _log(self._outputs, "")
                continue

            alloc_str = "  ".join(f"[dim]{n}×{c}[/dim]" for n, c in allocation.items())
            _log(self._outputs, f"  [bold]GENERATE[/bold]  {alloc_str}  model={gen_name}")
            _log(self._outputs, "")

            candidates: list[tuple[str, str]] = []

            for strategy_name, count in allocation.items():
                strategy = strategies.get(strategy_name)
                if strategy is None:
                    continue
                _status(
                    self._outputs,
                    f"Iter {iteration}/{self._iterations} · generating {count} {strategy_name} probes…",
                )
                gen_t0 = time.monotonic()
                try:
                    if strategy_name == "mutation":
                        prompts = strategy.generate(
                            pool.all_probes,
                            self._generator_model,
                            count=count,
                            technique_hints=technique_hints,
                            sample_fn=pool.sample_for_mutation,
                        )
                    else:
                        prompts = strategy.generate(
                            pool.all_probes,
                            self._generator_model,
                            count=count,
                            technique_hints=technique_hints,
                        )
                    candidates.extend((p, strategy_name) for p in prompts)
                    _log(
                        self._outputs,
                        f"  [dim]  ↳ {strategy_name}: {len(prompts)} probes  ({_elapsed(gen_t0)})[/dim]",
                    )
                except Exception as exc:
                    tb = traceback.format_exc()
                    _log(
                        self._outputs,
                        f"  [yellow]⚠ {strategy_name} generator error: {type(exc).__name__}: {exc}\n{tb}[/yellow]",
                    )

            _log(self._outputs, "")

            if not candidates:
                _log(self._outputs, "  [yellow]  no candidates — skipping probe loop[/yellow]")
                _log(self._outputs, "")
                all_results.extend(iter_results)
                continue

            for prompt, strategy_name in candidates:
                probe_index += 1

                _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, "")
                _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, "")

                _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()
                target_error: Exception | None = None
                content_filtered: bool = False
                try:
                    target_result = target_agent.run_sync(prompt)
                    response_text: str = target_result.output
                except Exception as exc:
                    target_error = exc
                    response_text = ""
                    if _is_content_filtered(exc):
                        content_filtered = True
                        _log(
                            self._outputs,
                            f"  [yellow]⚠ probe blocked by content filter ({tgt_name})[/yellow]",
                        )
                    else:
                        tb = traceback.format_exc()
                        _log(
                            self._outputs,
                            f"  [red]⚠ target error: {type(exc).__name__}: {exc}\n{tb}[/red]",
                        )
                tgt_elapsed = _elapsed(tgt_t0)

                _log(self._outputs, f"  [bold cyan]RESPONSE[/bold cyan]  [dim]{tgt_elapsed}[/dim]")
                _log(self._outputs, "  [dim]┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄┄[/dim]")
                if content_filtered:
                    _log(
                        self._outputs,
                        "  [yellow dim]  [FILTERED — content policy block][/yellow dim]",
                    )
                else:
                    for line in _block(_preview(response_text, _RESPONSE_PREVIEW)):
                        _log(self._outputs, line)
                _log(self._outputs, "")

                _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()
                if target_error is not None:
                    reason = (
                        "[FILTERED by content policy]"
                        if content_filtered
                        else f"[TARGET ERROR: {type(target_error).__name__}: {target_error}]"
                    )
                    goal_result = GoalResult(
                        achieved=False,
                        score=0.0,
                        reason=reason,
                        goal_name="objective",
                    )
                else:
                    try:
                        goal_result = judge.evaluate(prompt=prompt, response=response_text)
                    except Exception as exc:
                        tb = traceback.format_exc()
                        _log(
                            self._outputs,
                            f"  [red]⚠ judge error: {type(exc).__name__}: {exc}\n{tb}[/red]",
                        )
                        goal_result = GoalResult(
                            achieved=False,
                            score=0.0,
                            reason=f"[EVAL ERROR: {type(exc).__name__}: {exc}]",
                            goal_name="objective",
                        )
                jdg_elapsed = _elapsed(jdg_t0)

                bar = _score_bar(goal_result.score)
                verdict_tag = (
                    "[bold green]  ✓ ACHIEVED  [/bold green]"
                    if goal_result.achieved
                    else "[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, "")

                if goal_result.score >= _GRADING_SCORE_THRESHOLD:
                    _status(
                        self._outputs,
                        f"Iter {iteration}/{self._iterations} · probe {probe_index}/{total} · grading…",
                    )
                    _log(self._outputs, f"  [dim]🌱 grading seed  ({jdg_name})…[/dim]")
                    grade_t0 = time.monotonic()
                    analysis = grader.analyze(
                        prompt=prompt,
                        response=response_text,
                        existing_seeds=pool.graded,
                    )
                    if analysis is not None:
                        if grader.is_duplicate(analysis, pool.graded):
                            _log(
                                self._outputs,
                                f"  [dim]  ↩ duplicate technique (novelty={analysis.novelty:.2f}) — skipped  ({_elapsed(grade_t0)})[/dim]",
                            )
                        else:
                            new_probe = Probe(
                                title=f"[iter-{iteration}/{strategy_name}] {', '.join(analysis.techniques[:2]) or strategy_name}",
                                content=prompt,
                                category="generated",
                            )
                            graded_seed = GradedSeed(
                                probe=new_probe,
                                score=goal_result.score,
                                analysis=analysis,
                                iteration=iteration,
                            )
                            pool.add(graded_seed)
                            tags_str = "  ".join(f"[cyan]{t}[/cyan]" for t in analysis.techniques)
                            _log(
                                self._outputs,
                                f"  [green]  ✦ new seed added[/green]  novelty={analysis.novelty:.2f}  transfer={analysis.transferability:.2f}  ({_elapsed(grade_t0)})",
                            )
                            _log(self._outputs, f"    [dim]techniques: {tags_str}[/dim]")
                            _log(self._outputs, f"    [dim]{analysis.summary}[/dim]")
                    else:
                        _log(
                            self._outputs,
                            f"  [dim]  ⚠ grader returned no analysis  ({_elapsed(grade_t0)})[/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 pool.all_probes[:3]],
                    timestamp=datetime.datetime.utcnow().isoformat(),
                )
                for output in self._outputs:
                    output.emit(r)
                iter_results.append(r)

            all_results.extend(iter_results)

            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

            strategy_scores: dict[str, list[float]] = {}
            for r in iter_results:
                strategy_scores.setdefault(r.strategy_name, []).append(r.score)

            _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}")
            _log(self._outputs, f"    seed pool  {len(pool.graded)} graded seeds")

            for sname, scores in sorted(strategy_scores.items()):
                savg = sum(scores) / len(scores)
                _log(self._outputs, f"    [{sname}]  avg={savg:.2f}  n={len(scores)}")

            top_graded = pool.get_top_graded(3)
            if top_graded:
                _log(self._outputs, "    top seeds → next iter:")
                for gs in top_graded:
                    tags = ", ".join(gs.analysis.techniques[:3]) or "—"
                    _log(
                        self._outputs,
                        f"      [dim]{_score_bar(gs.weight, 6)} w={gs.weight:.2f}  [{tags}][/dim]",
                    )

            _log(self._outputs, "")

        # ── 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, f"  graded seeds  {len(pool.graded)}")

        top_techniques = pool.get_technique_histogram(top_n=5)
        if top_techniques:
            _log(self._outputs, f"  top techniques  {', '.join(top_techniques)}")

        all_strategy_scores: dict[str, list[float]] = {}
        for r in all_results:
            all_strategy_scores.setdefault(r.strategy_name, []).append(r.score)
        if all_strategy_scores:
            _log(self._outputs, "")
            _log(self._outputs, "  [bold]strategy breakdown:[/bold]")
            for sname, scores in sorted(
                all_strategy_scores.items(), key=lambda kv: -sum(kv[1]) / len(kv[1])
            ):
                savg = sum(scores) / len(scores)
                sbest = max(scores)
                n_ok = sum(1 for s in scores if s >= 0.5)
                _log(
                    self._outputs,
                    f"    {sname:<14}  n={len(scores)}  avg={savg:.2f}  best={sbest:.2f}  promising={n_ok}",
                )

        _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
