"""Pre-development deliberation engine.

Implements the multi-LLM, iterative deliberation pipeline that transforms a
seed prompt into PRD -> critique -> tech stack -> architecture -> implementation
plan through configurable rounds with peer exchange, primary decision-making,
and human-in-the-loop approval.

Usage:
    from abi_control_core.sdk.config_loader import load_config
    from abi_control_core.sdk.deliberation import DeliberationEngine

    cfg = load_config()
    engine = DeliberationEngine(cfg.pre_development, cfg.models)
    result = engine.run_pipeline(seed="Build a real-time dashboard")
"""

from __future__ import annotations

import json
import logging
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from datetime import UTC, datetime
from difflib import SequenceMatcher
from pathlib import Path
from typing import Any, Callable, Protocol

from .config_loader import (
    DeliberationPhaseConfig,
    ModelEntry,
    PreDevelopmentConfig,
)

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Types
# ---------------------------------------------------------------------------


class LLMInvoker(Protocol):
    """Protocol for invoking an LLM. Consumers must provide an implementation."""

    def invoke(
        self,
        model_entry: ModelEntry,
        prompt: str,
        max_tokens: int,
    ) -> LLMResponse: ...


@dataclass
class LLMResponse:
    """Response from an LLM invocation."""
    content: str
    model_id: str
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    latency_seconds: float = 0.0


@dataclass
class RoundArtifact:
    """Output from one participant in one round."""
    round_number: int
    participant_alias: str
    model_id: str
    content: str
    input_tokens: int = 0
    output_tokens: int = 0
    cost_usd: float = 0.0
    latency_seconds: float = 0.0
    timestamp: str = ""


@dataclass
class PhaseDecision:
    """The decision made after all rounds of a phase."""
    chosen_content: str
    decider_alias: str
    decider_model_id: str
    decision_policy: str
    cost_usd: float = 0.0


@dataclass
class HumanAction:
    """Records a human review action on a phase artifact."""
    action: str  # "approve", "reject", "modify"
    modified_content: str | None = None
    reviewer: str = ""
    timestamp: str = ""


@dataclass
class PhaseResult:
    """Complete result of one deliberation phase."""
    phase_name: str
    artifact_type: str
    rounds: list[list[RoundArtifact]] = field(default_factory=list)
    decision: PhaseDecision | None = None
    human_action: HumanAction | None = None
    final_content: str = ""
    total_cost_usd: float = 0.0
    total_input_tokens: int = 0
    total_output_tokens: int = 0
    wall_time_seconds: float = 0.0
    converged_early: bool = False
    skipped: bool = False


@dataclass
class PipelineResult:
    """Complete result of the full deliberation pipeline."""
    run_id: str = ""
    phases: list[PhaseResult] = field(default_factory=list)
    total_cost_usd: float = 0.0
    total_wall_time_seconds: float = 0.0
    seed_prompt: str = ""


# ---------------------------------------------------------------------------
# Default prompt templates
# ---------------------------------------------------------------------------

_PHASE_PROMPTS: dict[str, str] = {
    "seed_ideation": (
        "You are participating in a collaborative ideation session.\n"
        "Seed prompt: {seed}\n\n"
        "{peer_context}"
        "Produce a refined project concept with clear scope, goals, and constraints."
    ),
    "prd_generation": (
        "You are generating a Product Requirements Document (PRD).\n"
        "Project concept:\n{input_artifact}\n\n"
        "{peer_context}"
        "Produce a complete PRD with: problem statement, user stories, acceptance criteria, "
        "non-functional requirements, and success metrics."
    ),
    "prd_critique": (
        "You are critically reviewing a PRD for completeness, feasibility, and risks.\n"
        "PRD to review:\n{input_artifact}\n\n"
        "{peer_context}"
        "Produce a structured critique with: strengths, weaknesses, missing requirements, "
        "feasibility concerns, and specific improvement recommendations."
    ),
    "tech_stack": (
        "You are recommending a technology stack based on the PRD and architecture needs.\n"
        "PRD:\n{input_artifact}\n\n"
        "{peer_context}"
        "Produce a tech stack recommendation with: languages, frameworks, databases, "
        "infrastructure, and rationale for each choice."
    ),
    "architecture_plan": (
        "You are designing a software architecture.\n"
        "PRD:\n{input_artifact}\n"
        "Tech stack:\n{tech_stack}\n\n"
        "{peer_context}"
        "Produce an architecture plan with: component diagram description, data flow, "
        "API contracts, deployment topology, and key design decisions."
    ),
    "implementation_plan": (
        "You are creating a detailed implementation plan.\n"
        "PRD:\n{input_artifact}\n"
        "Architecture:\n{architecture}\n\n"
        "{peer_context}"
        "Produce an implementation plan with: phased milestones, task breakdown, "
        "dependency graph, estimated effort, risk mitigation, and testing strategy."
    ),
}


def _build_peer_context(peer_outputs: list[RoundArtifact]) -> str:
    if not peer_outputs:
        return ""
    lines = ["Previous round peer outputs:\n"]
    for art in peer_outputs:
        lines.append(f"--- {art.participant_alias} ({art.model_id}) ---")
        lines.append(art.content)
        lines.append("")
    lines.append("Consider the above when producing your response.\n\n")
    return "\n".join(lines)


def _similarity(a: str, b: str) -> float:
    """Compute text similarity ratio between two strings."""
    return SequenceMatcher(None, a, b).ratio()


# ---------------------------------------------------------------------------
# Engine
# ---------------------------------------------------------------------------


class DeliberationEngine:
    """Runs the pre-development deliberation pipeline.

    Args:
        config: PreDevelopmentConfig from the config loader.
        model_registry: Dict mapping model aliases to ModelEntry.
        invoker: LLM invoker implementation.
        human_callback: Optional callback for human review.
            Signature: (phase_name, artifact_content, human_review_mode) -> HumanAction
        artifact_dir: Directory for persisting artifacts.
    """

    def __init__(
        self,
        config: PreDevelopmentConfig,
        model_registry: dict[str, ModelEntry],
        invoker: LLMInvoker,
        human_callback: Callable[[str, str, str], HumanAction] | None = None,
        artifact_dir: str | Path | None = None,
    ):
        self.config = config
        self.models = model_registry
        self.invoker = invoker
        self.human_callback = human_callback
        self.artifact_dir = Path(artifact_dir or config.artifact_dir)

    def run_pipeline(
        self,
        seed: str,
        phases: list[str] | None = None,
    ) -> PipelineResult:
        """Run the full deliberation pipeline from seed to implementation plan.

        Args:
            seed: The initial seed prompt / project idea.
            phases: Ordered phase names to run. Defaults to the standard sequence.

        Returns:
            PipelineResult with all phase outputs and costs.
        """
        run_id = f"delib-{datetime.now(UTC).strftime('%Y%m%dT%H%M%S')}-{uuid.uuid4().hex[:8]}"
        result = PipelineResult(run_id=run_id, seed_prompt=seed)

        if not self.config.enabled:
            logger.info("Pre-development deliberation is disabled")
            return result

        pipeline_start = time.monotonic()
        total_cost = 0.0

        phase_order = phases or [
            "seed_ideation",
            "prd_generation",
            "prd_critique",
            "tech_stack",
            "architecture_plan",
            "implementation_plan",
        ]

        # Track artifacts from prior phases for context
        artifacts: dict[str, str] = {"seed": seed}

        for phase_name in phase_order:
            phase_cfg = self.config.phases.get(phase_name)
            if not phase_cfg or not phase_cfg.enabled:
                result.phases.append(PhaseResult(
                    phase_name=phase_name,
                    artifact_type=phase_name,
                    skipped=True,
                ))
                continue

            # Budget guard
            if total_cost >= self.config.global_max_cost_usd:
                logger.warning("Global cost limit reached (%.2f USD), skipping %s", total_cost, phase_name)
                result.phases.append(PhaseResult(
                    phase_name=phase_name,
                    artifact_type=phase_cfg.artifact_type,
                    skipped=True,
                ))
                continue

            elapsed = time.monotonic() - pipeline_start
            if elapsed >= self.config.global_max_wall_time_seconds:
                logger.warning("Global wall time limit reached (%.0fs), skipping %s", elapsed, phase_name)
                result.phases.append(PhaseResult(
                    phase_name=phase_name,
                    artifact_type=phase_cfg.artifact_type,
                    skipped=True,
                ))
                continue

            phase_result = self._run_phase(phase_name, phase_cfg, artifacts)
            result.phases.append(phase_result)
            total_cost += phase_result.total_cost_usd

            if phase_result.final_content:
                artifacts[phase_name] = phase_result.final_content

            # If human rejected, stop pipeline
            if phase_result.human_action and phase_result.human_action.action == "reject":
                logger.info("Human rejected %s — stopping pipeline", phase_name)
                break

        result.total_cost_usd = total_cost
        result.total_wall_time_seconds = time.monotonic() - pipeline_start

        # Persist pipeline result
        self._persist_result(result)

        return result

    def _run_phase(
        self,
        phase_name: str,
        phase_cfg: DeliberationPhaseConfig,
        prior_artifacts: dict[str, str],
    ) -> PhaseResult:
        """Run a single deliberation phase through all configured rounds."""
        phase_start = time.monotonic()
        phase_result = PhaseResult(
            phase_name=phase_name,
            artifact_type=phase_cfg.artifact_type,
        )

        # Resolve participant models
        participants: list[tuple[str, ModelEntry]] = []
        for alias in phase_cfg.participants:
            model = self.models.get(alias)
            if model:
                participants.append((alias, model))
            else:
                logger.warning("Model alias '%s' not in registry, skipping", alias)

        if not participants:
            logger.error("No valid participants for phase %s", phase_name)
            phase_result.skipped = True
            return phase_result

        prev_round_outputs: list[RoundArtifact] = []

        for round_num in range(1, phase_cfg.rounds + 1):
            # Wall time check
            elapsed = time.monotonic() - phase_start
            if elapsed >= phase_cfg.max_wall_time_seconds:
                logger.info("Phase %s hit wall time at round %d", phase_name, round_num)
                break

            # Cost check
            if phase_result.total_cost_usd >= phase_cfg.max_cost_usd:
                logger.info("Phase %s hit cost limit at round %d", phase_name, round_num)
                break

            round_artifacts: list[RoundArtifact] = []

            # Invoke participants with bounded concurrency (max_concurrency setting)
            max_workers = min(phase_cfg.max_concurrency, len(participants))
            invocations = [
                (alias, model, self._build_prompt(
                    phase_name, phase_cfg, prior_artifacts, prev_round_outputs, alias,
                ))
                for alias, model in participants
            ]

            if max_workers > 1 and len(invocations) > 1:
                with ThreadPoolExecutor(max_workers=max_workers) as pool:
                    futures = {
                        pool.submit(
                            self._invoke_one, alias, model, prompt,
                            round_num, phase_name, phase_cfg.max_tokens,
                        ): alias
                        for alias, model, prompt in invocations
                    }
                    for future in as_completed(futures):
                        art = future.result()
                        if art is not None:
                            round_artifacts.append(art)
                            phase_result.total_cost_usd += art.cost_usd
                            phase_result.total_input_tokens += art.input_tokens
                            phase_result.total_output_tokens += art.output_tokens
                # Preserve deterministic ordering by participant alias
                round_artifacts.sort(key=lambda a: a.participant_alias)
            else:
                for alias, model, prompt in invocations:
                    art = self._invoke_one(alias, model, prompt, round_num, phase_name, phase_cfg.max_tokens)
                    if art is not None:
                        round_artifacts.append(art)
                        phase_result.total_cost_usd += art.cost_usd
                        phase_result.total_input_tokens += art.input_tokens
                        phase_result.total_output_tokens += art.output_tokens

            phase_result.rounds.append(round_artifacts)

            # Convergence check
            if (
                phase_cfg.convergence.enabled
                and round_num >= phase_cfg.convergence.min_rounds
                and len(phase_result.rounds) >= 2
            ):
                prev = phase_result.rounds[-2]
                curr = phase_result.rounds[-1]
                if prev and curr:
                    # Compare primary decider output (or first participant)
                    prev_text = prev[0].content if prev else ""
                    curr_text = curr[0].content if curr else ""
                    sim = _similarity(prev_text, curr_text)
                    if sim >= phase_cfg.convergence.similarity_threshold:
                        logger.info(
                            "Phase %s converged at round %d (similarity=%.3f)",
                            phase_name, round_num, sim,
                        )
                        phase_result.converged_early = True
                        break

            # Update peer context for next round (all_to_all)
            prev_round_outputs = round_artifacts

        # Decision: primary decider selects from last round
        last_round = phase_result.rounds[-1] if phase_result.rounds else []
        if last_round:
            decider_alias = phase_cfg.primary_decider or participants[0][0]
            # Find decider's output
            chosen = None
            for art in last_round:
                if art.participant_alias == decider_alias:
                    chosen = art
                    break
            if not chosen:
                chosen = last_round[0]  # fallback to first

            phase_result.decision = PhaseDecision(
                chosen_content=chosen.content,
                decider_alias=chosen.participant_alias,
                decider_model_id=chosen.model_id,
                decision_policy=phase_cfg.decision_policy,
                cost_usd=phase_result.total_cost_usd,
            )
            phase_result.final_content = chosen.content

        # Human review
        if phase_cfg.human_review != "off" and self.human_callback and phase_result.final_content:
            human_action = self.human_callback(
                phase_name, phase_result.final_content, phase_cfg.human_review
            )
            phase_result.human_action = human_action

            if human_action.action == "modify" and human_action.modified_content:
                phase_result.final_content = human_action.modified_content
            elif human_action.action == "reject":
                phase_result.final_content = ""

        phase_result.wall_time_seconds = time.monotonic() - phase_start
        return phase_result

    def _invoke_one(
        self,
        alias: str,
        model: ModelEntry,
        prompt: str,
        round_num: int,
        phase_name: str,
        max_tokens: int = 8192,
    ) -> RoundArtifact | None:
        """Invoke a single participant and return its RoundArtifact (or None on error)."""
        invoke_start = time.monotonic()
        try:
            response = self.invoker.invoke(model, prompt, max_tokens)
        except Exception as exc:
            logger.error("LLM invocation failed for %s/%s: %s", phase_name, alias, exc)
            return None
        return RoundArtifact(
            round_number=round_num,
            participant_alias=alias,
            model_id=model.model_id,
            content=response.content,
            input_tokens=response.input_tokens,
            output_tokens=response.output_tokens,
            cost_usd=response.cost_usd,
            latency_seconds=time.monotonic() - invoke_start,
            timestamp=datetime.now(UTC).isoformat(),
        )

    def _build_prompt(
        self,
        phase_name: str,
        phase_cfg: DeliberationPhaseConfig,
        prior_artifacts: dict[str, str],
        peer_outputs: list[RoundArtifact],
        current_alias: str,
    ) -> str:
        """Build the prompt for a participant in a given round."""
        template = _PHASE_PROMPTS.get(phase_name, _PHASE_PROMPTS.get("seed_ideation", ""))

        # Filter peer outputs (all_to_all: show all; ring: show previous)
        if phase_cfg.peer_exchange_mode == "none":
            peer_context = ""
        else:
            # all_to_all: show all peer outputs except own
            relevant_peers = [p for p in peer_outputs if p.participant_alias != current_alias]
            peer_context = _build_peer_context(relevant_peers)

        return template.format(
            seed=prior_artifacts.get("seed", ""),
            input_artifact=prior_artifacts.get(
                self._input_artifact_key(phase_name), prior_artifacts.get("seed", "")
            ),
            tech_stack=prior_artifacts.get("tech_stack", ""),
            architecture=prior_artifacts.get("architecture_plan", ""),
            peer_context=peer_context,
        )

    def _input_artifact_key(self, phase_name: str) -> str:
        """Determine which prior artifact feeds into this phase."""
        mapping = {
            "seed_ideation": "seed",
            "prd_generation": "seed_ideation",
            "prd_critique": "prd_generation",
            "tech_stack": "prd_generation",
            "architecture_plan": "prd_generation",
            "implementation_plan": "architecture_plan",
        }
        return mapping.get(phase_name, "seed")

    def _persist_result(self, result: PipelineResult) -> None:
        """Write pipeline result artifacts to disk."""
        try:
            run_dir = self.artifact_dir / result.run_id
            run_dir.mkdir(parents=True, exist_ok=True)

            # Summary JSON
            summary = {
                "run_id": result.run_id,
                "seed_prompt": result.seed_prompt,
                "total_cost_usd": result.total_cost_usd,
                "total_wall_time_seconds": result.total_wall_time_seconds,
                "phases": [],
            }

            for phase in result.phases:
                phase_data: dict[str, Any] = {
                    "phase_name": phase.phase_name,
                    "artifact_type": phase.artifact_type,
                    "skipped": phase.skipped,
                    "converged_early": phase.converged_early,
                    "total_cost_usd": phase.total_cost_usd,
                    "total_input_tokens": phase.total_input_tokens,
                    "total_output_tokens": phase.total_output_tokens,
                    "wall_time_seconds": phase.wall_time_seconds,
                    "rounds_count": len(phase.rounds),
                }

                if phase.decision:
                    phase_data["decision"] = {
                        "decider_alias": phase.decision.decider_alias,
                        "decider_model_id": phase.decision.decider_model_id,
                        "decision_policy": phase.decision.decision_policy,
                    }

                if phase.human_action:
                    phase_data["human_action"] = {
                        "action": phase.human_action.action,
                        "reviewer": phase.human_action.reviewer,
                        "timestamp": phase.human_action.timestamp,
                    }

                summary["phases"].append(phase_data)

                # Write per-phase artifacts
                if phase.final_content:
                    phase_file = run_dir / f"{phase.phase_name}.md"
                    phase_file.write_text(phase.final_content, encoding="utf-8")

                # Write round details
                for round_idx, round_arts in enumerate(phase.rounds):
                    for art in round_arts:
                        art_file = run_dir / f"{phase.phase_name}_r{round_idx + 1}_{art.participant_alias}.md"
                        art_file.write_text(art.content, encoding="utf-8")

            summary_file = run_dir / "pipeline_summary.json"
            summary_file.write_text(
                json.dumps(summary, indent=2, default=str), encoding="utf-8"
            )

            logger.info("Deliberation artifacts written to %s", run_dir)
        except Exception as exc:
            logger.warning("Failed to persist deliberation artifacts: %s", exc)
