"""Review phase for deliberate."""

import asyncio
import json
from dataclasses import dataclass

from deliberate.adapters.base import ModelAdapter
from deliberate.budget.tracker import BudgetTracker
from deliberate.prompts.review import CRITERIA_DESCRIPTIONS, REVIEW_PROMPT
from deliberate.types import ExecutionResult, Review, Score, VoteResult
from deliberate.voting.aggregation import (
    aggregate_approval,
    aggregate_borda,
    aggregate_weighted_borda,
    calculate_confidence,
    get_rankings,
)


@dataclass
class ReviewPhase:
    """Orchestrates the review phase of the jury workflow.

    Multiple agents review execution results and vote on
    the best candidate using configured aggregation.
    """

    agents: list[str]
    adapters: dict[str, ModelAdapter]
    budget: BudgetTracker
    criteria: list[str]
    scale: str = "1-10"
    aggregation_method: str = "borda"  # borda | approval | weighted_borda
    approval_threshold: float = 0.7

    async def run(
        self,
        task: str,
        candidates: list[ExecutionResult],
    ) -> tuple[list[Review], VoteResult | None]:
        """Run the review phase.

        Args:
            task: The original task description.
            candidates: List of execution results to review.

        Returns:
            Tuple of (reviews, vote_result).
        """
        # Filter to successful candidates with diffs
        valid = [c for c in candidates if c.success and c.diff]

        if not valid:
            return [], None

        reviews = await self._collect_reviews(task, valid)

        if not reviews:
            # No reviews collected, return first candidate as winner
            return [], VoteResult(
                winner_id=valid[0].id,
                rankings=[c.id for c in valid],
                scores={c.id: 0.5 for c in valid},
                vote_breakdown={},
                confidence=0.0,
            )

        vote_result = self._aggregate(valid, reviews)
        return reviews, vote_result

    async def _collect_reviews(
        self,
        task: str,
        candidates: list[ExecutionResult],
    ) -> list[Review]:
        """Collect reviews from all reviewers for all candidates."""

        async def review_one(
            reviewer: str,
            candidate: ExecutionResult,
        ) -> Review | None:
            # Don't allow self-review
            if reviewer == candidate.agent:
                return None

            adapter = self.adapters.get(reviewer)
            if not adapter:
                return None

            # Build criteria text
            criteria_text = "\n".join(
                f"- {c}: {CRITERIA_DESCRIPTIONS.get(c, 'Evaluate this criterion')}"
                for c in self.criteria
            )

            prompt = REVIEW_PROMPT.format(
                task=task,
                summary=candidate.summary,
                diff=candidate.diff or "(no diff)",
                criteria=criteria_text,
                scale=self.scale,
            )

            try:
                response = await adapter.call(prompt)
                self.budget.record_usage(reviewer, response.token_usage)
                return self._parse_review(
                    reviewer,
                    candidate.id,
                    response.content,
                    response.token_usage,
                )
            except Exception as e:
                print(f"Warning: {reviewer} review of {candidate.id} failed: {e}")
                return None

        # Create all review tasks
        tasks = [review_one(r, c) for r in self.agents for c in candidates]
        results = await asyncio.gather(*tasks)
        return [r for r in results if r is not None]

    def _parse_review(
        self,
        reviewer: str,
        candidate_id: str,
        response: str,
        tokens: int,
    ) -> Review:
        """Parse a review response into a Review object."""
        try:
            # Try to extract JSON from the response
            # Look for JSON block in markdown
            json_match = response
            if "```json" in response:
                start = response.find("```json") + 7
                end = response.find("```", start)
                if end > start:
                    json_match = response[start:end]
            elif "```" in response:
                start = response.find("```") + 3
                end = response.find("```", start)
                if end > start:
                    json_match = response[start:end]

            data = json.loads(json_match.strip())

            # Parse scores
            raw_scores = data.get("scores", {})
            scores = []
            for criterion in self.criteria:
                raw_value = float(raw_scores.get(criterion, 5))
                # Normalize to 0-1 (assuming 1-10 scale)
                normalized = raw_value / 10.0
                scores.append(
                    Score(
                        criterion=criterion,
                        value=normalized,
                        raw_value=raw_value,
                    )
                )

            overall = float(data.get("overall", 5)) / 10.0

            return Review(
                reviewer=reviewer,
                candidate_id=candidate_id,
                scores=scores,
                overall_score=overall,
                recommendation=data.get("recommendation", "accept"),
                comments=data.get("reasoning"),
                token_usage=tokens,
            )

        except (json.JSONDecodeError, KeyError, ValueError):
            # If parsing fails, return a neutral review
            return Review(
                reviewer=reviewer,
                candidate_id=candidate_id,
                scores=[Score(c, 0.5, 5) for c in self.criteria],
                overall_score=0.5,
                recommendation="accept",
                comments=response[:500] if response else None,
                token_usage=tokens,
            )

    def _aggregate(
        self,
        candidates: list[ExecutionResult],
        reviews: list[Review],
    ) -> VoteResult:
        """Aggregate reviews into a voting result."""
        # Build vote breakdown: reviewer -> candidate -> score
        vote_breakdown: dict[str, dict[str, float]] = {}
        for review in reviews:
            if review.reviewer not in vote_breakdown:
                vote_breakdown[review.reviewer] = {}
            vote_breakdown[review.reviewer][review.candidate_id] = review.overall_score

        candidate_ids = [c.id for c in candidates]

        # Apply the configured aggregation method
        if self.aggregation_method == "approval":
            scores = aggregate_approval(
                candidate_ids,
                vote_breakdown,
                threshold=self.approval_threshold,
            )
        elif self.aggregation_method == "weighted_borda":
            scores = aggregate_weighted_borda(candidate_ids, vote_breakdown)
        else:  # default to borda
            scores = aggregate_borda(candidate_ids, vote_breakdown)

        rankings = get_rankings(scores)
        confidence = calculate_confidence(scores)

        return VoteResult(
            winner_id=rankings[0] if rankings else candidate_ids[0],
            rankings=rankings,
            scores=scores,
            vote_breakdown=vote_breakdown,
            confidence=confidence,
        )
