"""Fact Checker Agent

Simple agent that verifies review claims using codebase exploration.
Returns indices of reviews to keep - caller handles filtering.
"""

from __future__ import annotations

import json
import logging
from typing import TYPE_CHECKING, Annotated

from pydantic import BaseModel, ConfigDict, Field

from ..schema import (
    BasicReviewOutput,
    GitHubReviewComment,
    GitLabReviewComment,
    MergeRequest,
)
from .base import AgentConfig, AgentResponse, BaseAgentImpl, ReasoningEffort, TokenUsage
from .tools import build_codebase_tools, make_tool_handler
from .utils import build_model_name, format_template

if TYPE_CHECKING:
    from ..explorer import CodeExplorer
    from ..workflows.context import LinkedRepoInfo, PipelineContext

logger = logging.getLogger(__name__)

# Type alias for platform review types
PlatformReviewType = GitHubReviewComment | GitLabReviewComment | BasicReviewOutput


def _extract_review_location(
    review: GitHubReviewComment | GitLabReviewComment,
) -> dict[str, str | int | None]:
    """Extract path and line from a platform-specific review."""
    if isinstance(review, GitHubReviewComment):
        return {"path": review.path, "line": review.line}
    elif isinstance(review, GitLabReviewComment):
        return {"path": review.new_path, "line": review.new_line}
    return {"path": None, "line": None}


class ReasonEntry(BaseModel):
    """A single reason for a fact check decision."""

    model_config = ConfigDict(extra="forbid")

    index: int = Field(description="The review index (0-based)")
    reason: str = Field(description="The reason: 'KEEP: ...' or 'DISCARD: ...'")


class FactCheckResult(BaseModel):
    """Result of fact checking - just indices to keep."""

    model_config = ConfigDict(extra="forbid")

    keep_indices: Annotated[
        list[int], Field(description="Indices of reviews that are VERIFIED as valid bugs")
    ] = []


class FactCheckResultDebug(BaseModel):
    """Result of fact checking with reasons (debug mode)."""

    model_config = ConfigDict(extra="forbid")

    keep_indices: Annotated[
        list[int], Field(description="Indices of reviews that are VERIFIED as valid bugs")
    ] = []
    reasons: list[ReasonEntry] = Field(
        default_factory=list,
        description="Reason for each review decision",
    )


class FactCheckerAgent:
    """Simple fact checker with exploration tool.

    Flow:
    1. Receives reviews to verify (in batches of BATCH_SIZE)
    2. Can call explore() to check claims against codebase
    3. Returns indices of reviews to keep
    """

    MAX_ITERATIONS = 25
    MAX_TOOL_CALLS = 20
    BATCH_SIZE = 10

    def __init__(
        self,
        config: Annotated[AgentConfig, Field(description="LLM configuration")],
        reasoning_effort: ReasoningEffort = None,
    ):
        self.base = BaseAgentImpl(
            name="FactChecker",
            description="Verifies review claims using codebase exploration",
            config=config,
            reasoning_effort=reasoning_effort,
        )

    async def _verify_batch(
        self,
        merge_info: MergeRequest,
        diffs: str,
        batch: list[tuple[int, PlatformReviewType]],
        code_explorer: CodeExplorer,
        batch_num: int = 0,
        total_batches: int = 1,
        user_guidelines: str = "",
        debug: bool = False,
        linked_repos: list[LinkedRepoInfo] | None = None,
        team_guidelines: str | None = None,
    ) -> AgentResponse[FactCheckResult | FactCheckResultDebug]:
        """Verify a single batch of reviews.

        Args:
            merge_info: The merge request information
            diffs: The code diffs
            batch: List of (original_index, review) tuples
            code_explorer: CodeExplorer for codebase exploration
            batch_num: Batch number (0-indexed)
            total_batches: Total number of batches
            user_guidelines: Optional user guidelines
            debug: If True, include reasons for each decision
            linked_repos: Optional list of linked repository info for prompt rendering

        Returns:
            AgentResponse with keep_indices mapped to ORIGINAL indices
        """
        model_name = build_model_name(self.base.config.provider, self.base.config.model)
        batch_label = f"[Batch {batch_num + 1}/{total_batches}]"

        if debug:
            review_indices = [orig_idx for orig_idx, _ in batch]
            logger.info(f"{batch_label} Processing reviews: {review_indices}")
        response_model = FactCheckResultDebug if debug else FactCheckResult

        # Format reviews for the prompt (use batch-local indices 0, 1, ...)
        reviews_list = []
        for batch_idx, (_, r) in enumerate(batch):
            if isinstance(r, (GitHubReviewComment, GitLabReviewComment)):
                loc = _extract_review_location(r)
                reviews_list.append(
                    {"index": batch_idx, "path": loc["path"], "line": loc["line"], "body": r.body}
                )
            else:
                reviews_list.append(
                    {"index": batch_idx, "path": None, "line": None, "body": r.body}
                )
        reviews_json = json.dumps(reviews_list, indent=2)

        # Extract changed file paths from the diff for scoped searches
        diff_files = [line[6:] for line in diffs.splitlines() if line.startswith("+++ b/")]

        # Build context for the prompt
        ctx = {
            "title": merge_info.title,
            "description": merge_info.description,
            "target_branch": merge_info.target_branch,
            "source_branch": merge_info.source_branch,
            "project_id": merge_info.repo,
            "reviews": reviews_json,
            "kg_available": True,
            "user_guidelines": user_guidelines,
            "debug": debug,
            "linked_repos": linked_repos or [],
            "team_guidelines": team_guidelines or "",
            "diff_files": diff_files,
            "max_tool_calls": self.MAX_TOOL_CALLS,
        }

        # Build project names for grep enum when linked repos exist
        project_names: list[str] | None = None
        if linked_repos:
            project_names = ["primary"] + [r["name"] for r in linked_repos]
        tools = build_codebase_tools(project_names)

        # Load template and format
        template = self.base.load_prompt("fact_checker_agent.txt")
        system_prompt = format_template(self.base.jinja_env, template, ctx)
        output_format = self.base.get_output_format(response_model)

        try:
            handle_tool = make_tool_handler(
                code_explorer,
                f"FactChecker {batch_label}",
                self.MAX_TOOL_CALLS,
            )

            result = await self.base.analyze(
                system_prompt=system_prompt,
                user_prompt=output_format,
                response_model=response_model,
                cached_content=diffs,
                tools=tools,
                tool_handler=handle_tool,
                max_iterations=self.MAX_ITERATIONS,
            )

            # Enforce tool usage: retry if zero tool calls
            tool_calls_used = handle_tool.state["count"]  # type: ignore[attr-defined]
            if tool_calls_used == 0:
                logger.warning(
                    f"FactChecker {batch_label} made ZERO tool calls. Retrying with enforcement..."
                )
                handle_tool = make_tool_handler(
                    code_explorer,
                    f"FactChecker {batch_label}",
                    self.MAX_TOOL_CALLS,
                )
                enforced_prompt = (
                    "REJECTED: Your previous response used NO tools. "
                    "You MUST call grep() and/or read_lines() to verify each review against the actual codebase. "
                    "Decisions made without searching the code are invalid. Search now, then decide.\n\n"
                    + output_format
                )
                result = await self.base.analyze(
                    system_prompt=system_prompt,
                    user_prompt=enforced_prompt,
                    response_model=response_model,
                    cached_content=diffs,
                    tools=tools,
                    tool_handler=handle_tool,
                    max_iterations=self.MAX_ITERATIONS,
                )
                tool_calls_used = handle_tool.state["count"]  # type: ignore[attr-defined]
                if tool_calls_used == 0:
                    logger.error(
                        f"FactChecker {batch_label} STILL made ZERO tool calls after retry!"
                    )

            logger.info(
                f"FactChecker {batch_label} finished: {handle_tool.state['count']} tool calls"  # type: ignore[attr-defined]
            )

            # Map batch-local indices back to original indices
            index_map = {
                batch_idx: original_idx for batch_idx, (original_idx, _) in enumerate(batch)
            }
            mapped_keep_indices = [
                index_map[i] for i in result.output.keep_indices if i in index_map
            ]

            if debug:
                mapped_reasons = [
                    ReasonEntry(index=index_map.get(r.index, r.index), reason=r.reason)
                    for r in result.output.reasons
                    if r.index in index_map
                ]
                return AgentResponse(
                    output=FactCheckResultDebug(
                        keep_indices=mapped_keep_indices, reasons=mapped_reasons
                    ),
                    metadata=result.metadata,
                )
            else:
                return AgentResponse(
                    output=FactCheckResult(keep_indices=mapped_keep_indices),
                    metadata=result.metadata,
                )

        except Exception as e:
            logger.error(f"FactChecker batch failed: {e}")
            # On failure, keep all reviews in this batch (fail open)
            return AgentResponse(
                output=FactCheckResult(keep_indices=[orig_idx for orig_idx, _ in batch]),
                metadata=TokenUsage.zero(model=model_name),
            )

    async def verify(
        self,
        merge_info: MergeRequest,
        diffs: str,
        reviews: list[PlatformReviewType],
        code_explorer: CodeExplorer,
        user_guidelines: str = "",
        debug: bool = False,
        linked_repos: list[LinkedRepoInfo] | None = None,
        team_guidelines: str | None = None,
    ) -> AgentResponse[FactCheckResult | FactCheckResultDebug]:
        """Verify reviews and return indices to keep.

        Reviews are processed in batches of BATCH_SIZE concurrently.

        Args:
            merge_info: The merge request information
            diffs: The code diffs
            reviews: AI-generated reviews to fact-check
            code_explorer: CodeExplorer for codebase exploration
            user_guidelines: Optional user guidelines
            debug: If True, include reasons for each decision
            linked_repos: Optional list of linked repository info for prompt rendering

        Returns:
            AgentResponse with FactCheckResult containing keep_indices (and reasons if debug)
        """
        import asyncio

        model_name = build_model_name(self.base.config.provider, self.base.config.model)

        if not reviews:
            return AgentResponse(
                output=FactCheckResult(keep_indices=[]),
                metadata=TokenUsage.zero(model=model_name),
            )

        # Debug: log context for test dataset
        if debug:
            logger.info(f"[FACT_CHECK_CONTEXT] repo={merge_info.repo} pr={merge_info.id}")
            logger.info(f"[FACT_CHECK_CONTEXT] title={merge_info.title}")
            logger.info(f"[FACT_CHECK_CONTEXT] diff_lines={len(diffs.splitlines())}")

        # Create batches of (original_index, review) tuples
        indexed_reviews = list(enumerate(reviews))
        batches = [
            indexed_reviews[i : i + self.BATCH_SIZE]
            for i in range(0, len(indexed_reviews), self.BATCH_SIZE)
        ]

        logger.info(f"FactChecker: processing {len(reviews)} reviews in {len(batches)} batches")

        # Debug: log input reviews for test dataset
        if debug:
            logger.info("[FACT_CHECK_INPUT] Reviews to verify:")
            for i, review in enumerate(reviews):
                logger.info(f"[REVIEW_{i}] {review.body}")
            for i, batch in enumerate(batches):
                review_indices = [orig_idx for orig_idx, _ in batch]
                logger.info(f"  Batch {i + 1}/{len(batches)}: reviews {review_indices}")

        # Process all batches concurrently
        # NOTE: Each batch has its own conversation - exploration results are NOT shared
        batch_tasks = [
            self._verify_batch(
                merge_info=merge_info,
                diffs=diffs,
                batch=batch,
                code_explorer=code_explorer,
                batch_num=i,
                total_batches=len(batches),
                user_guidelines=user_guidelines,
                debug=debug,
                linked_repos=linked_repos,
                team_guidelines=team_guidelines,
            )
            for i, batch in enumerate(batches)
        ]
        batch_results = await asyncio.gather(*batch_tasks)

        # Aggregate results
        all_keep_indices: list[int] = []
        all_reasons: list[ReasonEntry] = []
        total_usage = TokenUsage.zero(model=model_name)

        for result in batch_results:
            all_keep_indices.extend(result.output.keep_indices)
            total_usage = total_usage + result.metadata
            if debug and hasattr(result.output, "reasons"):
                all_reasons.extend(result.output.reasons)

        # Sort indices for consistent output
        all_keep_indices.sort()

        logger.info(
            f"FactChecker: kept {len(all_keep_indices)} out of {len(reviews)}, "
            f"tokens={total_usage.total_tokens} (cached={total_usage.cached_tokens})"
        )

        for i in range(len(reviews)):
            if i in all_keep_indices:
                logger.info(f"  KEEP [{i}]")
            else:
                logger.info(f"  DISCARD [{i}]")

        # Debug: log decisions with reasons for test dataset
        if debug:
            logger.info("[FACT_CHECK_OUTPUT] Decisions:")
            for reason in all_reasons:
                logger.info(f"[DECISION_{reason.index}] {reason.reason}")
            return AgentResponse(
                output=FactCheckResultDebug(keep_indices=all_keep_indices, reasons=all_reasons),
                metadata=total_usage,
            )
        else:
            return AgentResponse(
                output=FactCheckResult(keep_indices=all_keep_indices),
                metadata=total_usage,
            )

    async def run(
        self,
        ctx: PipelineContext,
    ) -> None:
        """Execute as a pipeline step.

        Verifies ctx.reviews and filters to only verified ones.
        """
        if ctx.merge_request is None:
            raise ValueError("merge_request must be set before running fact checker")
        if not ctx.reviews:
            logger.info("FactChecker skipped: no reviews")
            return
        if ctx.code_explorer is None:
            raise ValueError("code_explorer must be set before running fact checker")

        logger.info(f"Starting FactChecker on {len(ctx.reviews)} reviews")
        before_count = len(ctx.reviews)

        self.base.reasoning_effort = ctx.reasoning_effort

        result = await self.verify(
            merge_info=ctx.merge_request,
            diffs=ctx.merge_request.diffs,
            reviews=ctx.reviews,
            code_explorer=ctx.code_explorer,
            user_guidelines=ctx.user_guidelines,
            debug=ctx.debug,
            linked_repos=ctx.linked_repos,
            team_guidelines=ctx.team_guidelines,
        )

        # Filter reviews to only those that passed verification
        ctx.reviews = [ctx.reviews[i] for i in result.output.keep_indices if i < len(ctx.reviews)]
        ctx.track("fact_checker", result.metadata)
        logger.info(f"FactChecker completed: {before_count} -> {len(ctx.reviews)} reviews")

    @classmethod
    def default(cls) -> FactCheckerAgent:
        """Create FactCheckerAgent with default config from environment."""
        from ..config import Config

        config = Config.from_env()
        return cls(config.create_agent_config("fact_checker"))
