"""Core scan engine — orchestrates attacks against a target."""

from __future__ import annotations

import asyncio
import time
from datetime import datetime, timezone

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

from llm_pentest.analyzer import analyze_response, analyze_with_judge
from llm_pentest.connector import TargetConnector
from llm_pentest.judge import LLMJudge
from llm_pentest.models import AttackResult, Chain, Payload, ScanResult, TargetConfig, TurnPrompt

console = Console()


async def _analyze(judge: LLMJudge | None, payload: Payload, prompt: str, response: str, ms: float) -> AttackResult:
    if judge:
        return await analyze_with_judge(judge, payload, prompt, response, ms)
    return analyze_response(payload, prompt, response, ms)


async def _run_single(
    connector: TargetConnector, payload: Payload, prompt: str, dry_run: bool, judge: LLMJudge | None,
    sem: asyncio.Semaphore | None = None,
) -> AttackResult:
    if sem:
        async with sem:
            return await _do_single(connector, payload, prompt, dry_run, judge)
    return await _do_single(connector, payload, prompt, dry_run, judge)


async def _do_single(
    connector: TargetConnector, payload: Payload, prompt: str, dry_run: bool, judge: LLMJudge | None,
) -> AttackResult:
    if dry_run:
        return AttackResult(payload_id=payload.id, prompt=prompt, response="[DRY RUN]",
                            severity=payload.severity, owasp=payload.owasp)
    t0 = time.monotonic()
    try:
        response = await connector.send(prompt)
        ms = (time.monotonic() - t0) * 1000
        return await _analyze(judge, payload, prompt, response, ms)
    except Exception as e:
        return AttackResult(
            payload_id=payload.id, prompt=prompt, response="",
            error=str(e), duration_ms=(time.monotonic() - t0) * 1000,
            severity=payload.severity, owasp=payload.owasp,
        )


async def _run_multi_turn(
    connector: TargetConnector, payload: Payload, turns: list[TurnPrompt], dry_run: bool,
    judge: LLMJudge | None, sem: asyncio.Semaphore | None = None,
) -> AttackResult:
    if sem:
        async with sem:
            return await _do_multi_turn(connector, payload, turns, dry_run, judge)
    return await _do_multi_turn(connector, payload, turns, dry_run, judge)


async def _do_multi_turn(
    connector: TargetConnector, payload: Payload, turns: list[TurnPrompt], dry_run: bool,
    judge: LLMJudge | None,
) -> AttackResult:
    history: list[dict[str, str]] = []
    last_response = ""
    total_ms = 0.0
    prompts_log: list[str] = []

    for turn in turns:
        prompts_log.append(f"[{turn.role}] {turn.content}")
        if dry_run:
            history.append({"role": turn.role, "content": turn.content})
            history.append({"role": "assistant", "content": "[DRY RUN]"})
            continue
        t0 = time.monotonic()
        try:
            last_response = await connector.send(turn.content, history=history)
            total_ms += (time.monotonic() - t0) * 1000
        except Exception as e:
            return AttackResult(
                payload_id=payload.id, prompt="\n".join(prompts_log), response="",
                error=str(e), duration_ms=total_ms, severity=payload.severity, owasp=payload.owasp,
            )
        history.append({"role": turn.role, "content": turn.content})
        history.append({"role": "assistant", "content": last_response})

    combined_prompt = "\n".join(prompts_log)
    if dry_run:
        return AttackResult(payload_id=payload.id, prompt=combined_prompt, response="[DRY RUN]",
                            severity=payload.severity, owasp=payload.owasp)
    return await _analyze(judge, payload, combined_prompt, last_response, total_ms)


async def _run_chain(
    connector: TargetConnector, chain: Chain, payloads_by_id: dict[str, Payload],
    dry_run: bool, judge: LLMJudge | None, sem: asyncio.Semaphore | None = None,
) -> AttackResult:
    if sem:
        async with sem:
            return await _do_chain(connector, chain, payloads_by_id, dry_run, judge)
    return await _do_chain(connector, chain, payloads_by_id, dry_run, judge)


async def _do_chain(
    connector: TargetConnector, chain: Chain, payloads_by_id: dict[str, Payload],
    dry_run: bool, judge: LLMJudge | None,
) -> AttackResult:
    history: list[dict[str, str]] = []
    last_response = ""
    total_ms = 0.0
    prompts_log: list[str] = []
    last_payload: Payload | None = None

    for step in chain.steps:
        payload = payloads_by_id.get(step.payload_id)
        if not payload or step.prompt_index >= len(payload.prompts):
            continue
        last_payload = payload
        prompt = payload.prompts[step.prompt_index]
        prompts_log.append(f"[{step.payload_id}] {prompt}")

        if dry_run:
            history.append({"role": "user", "content": prompt})
            history.append({"role": "assistant", "content": "[DRY RUN]"})
            continue

        t0 = time.monotonic()
        try:
            last_response = await connector.send(prompt, history=history)
            total_ms += (time.monotonic() - t0) * 1000
        except Exception as e:
            return AttackResult(
                payload_id=chain.id, prompt="\n".join(prompts_log), response="",
                error=str(e), duration_ms=total_ms,
            )
        history.append({"role": "user", "content": prompt})
        history.append({"role": "assistant", "content": last_response})

    combined = "\n".join(prompts_log)
    if dry_run or not last_payload:
        return AttackResult(payload_id=chain.id, prompt=combined, response="[DRY RUN]")
    return await _analyze(judge, last_payload, combined, last_response, total_ms)


async def run_scan(
    config: TargetConfig,
    payloads: list[Payload],
    dry_run: bool = False,
    judge: LLMJudge | None = None,
    concurrency: int = 1,
    chains: list[Chain] | None = None,
) -> ScanResult:
    connector = TargetConnector(config)
    sem = asyncio.Semaphore(concurrency) if concurrency > 1 else None
    payloads_by_id = {p.id: p for p in payloads}

    tasks: list[asyncio.Task[AttackResult]] = []

    for payload in payloads:
        for prompt in payload.prompts:
            tasks.append(asyncio.ensure_future(
                _run_single(connector, payload, prompt, dry_run, judge, sem)))
        for turn_seq in payload.turns:
            tasks.append(asyncio.ensure_future(
                _run_multi_turn(connector, payload, turn_seq, dry_run, judge, sem)))

    for chain in (chains or []):
        tasks.append(asyncio.ensure_future(
            _run_chain(connector, chain, payloads_by_id, dry_run, judge, sem)))

    results: list[AttackResult] = []
    with Progress(SpinnerColumn(), TextColumn("{task.description}"), console=console) as progress:
        ptask = progress.add_task(f"Scanning ({len(tasks)} attacks)...", total=len(tasks))
        for coro in asyncio.as_completed(tasks):
            result = await coro
            results.append(result)
            progress.update(ptask, advance=1,
                            description=f"[cyan]{result.payload_id}[/] → {'✓' if result.success else '✗'}")

    successful = sum(1 for r in results if r.success)
    return ScanResult(
        target=config.name,
        timestamp=datetime.now(timezone.utc).isoformat(),
        total_attacks=len(results),
        successful_attacks=successful,
        results=results,
    )
