"""Plan-and-Solve Agent."""

import time
from dataclasses import dataclass
from typing import Any, Callable, Dict, Optional

from .planner import Planner
from .solver import Solver
from .checker import SelfChecker
from .types import Plan, Solution, CheckResult, PSTrajectory


@dataclass
class PSResult:
    """Result of a Plan-and-Solve session."""
    success: bool
    answer: str
    plan: Optional[Plan]
    solution: Optional[Solution]
    check_result: Optional[CheckResult]
    trajectory: PSTrajectory

    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "answer": self.answer,
            "plan": self.plan.to_dict() if self.plan else None,
            "solution": self.solution.to_dict() if self.solution else None,
            "check_result": self.check_result.to_dict() if self.check_result else None,
            "trajectory": self.trajectory.to_dict(),
        }


class PlanAndSolveAgent:
    """Agent that plans before solving.

    Plan-and-Solve improves zero-shot reasoning by explicitly
    generating a plan before execution, reducing reasoning errors.

    Example:
        >>> def plan_fn(problem):
        ...     return ["Step 1", "Step 2", "Step 3"]
        >>>
        >>> def solve_fn(step, context):
        ...     return f"Executed {step}"
        >>>
        >>> agent = PlanAndSolveAgent(plan_fn, solve_fn)
        >>> result = agent.solve("What is 15 * 17?")
    """

    def __init__(
        self,
        plan_fn: Optional[Callable[[str], list]] = None,
        solve_fn: Optional[Callable[[str, Dict], str]] = None,
        check_fn: Optional[Callable[[str, str, str], CheckResult]] = None,
        max_iterations: int = 3,
        verbose: bool = False,
    ):
        """Initialize the Plan-and-Solve agent.

        Args:
            plan_fn: Function to generate plan steps
            solve_fn: Function to execute steps
            check_fn: Function to validate solutions
            max_iterations: Maximum refinement iterations
            verbose: Whether to print debug info
        """
        self._planner = Planner(plan_fn=plan_fn, verbose=verbose)
        self._solver = Solver(solve_fn=solve_fn, verbose=verbose)
        self._checker = SelfChecker(check_fn=check_fn, verbose=verbose)
        self._max_iterations = max_iterations
        self._verbose = verbose

    def solve(self, problem: str) -> PSResult:
        """Solve a problem using Plan-and-Solve.

        Args:
            problem: Problem description

        Returns:
            PSResult with solution
        """
        start_time = time.time()
        trajectory = PSTrajectory(problem=problem)

        if self._verbose:
            print(f"\n{'='*50}")
            print(f"Plan-and-Solve: {problem[:50]}...")
            print(f"{'='*50}")

        for iteration in range(self._max_iterations):
            if self._verbose:
                print(f"\n--- Iteration {iteration + 1} ---")

            # Step 1: Plan
            plan = self._planner.create_plan(problem)
            trajectory.plan = plan

            if self._verbose:
                print(plan.to_prompt())

            # Step 2: Execute
            solution = self._solver.execute_plan(plan)
            trajectory.solution = solution

            # Step 3: Check
            check_result = self._checker.check(problem, plan, solution)
            trajectory.check_result = check_result

            if check_result.passed:
                trajectory.success = True
                break

            # Refine if needed
            if iteration < self._max_iterations - 1:
                if self._verbose:
                    print("\n[Agent] Refining based on check result...")

                # Update problem with feedback for next iteration
                problem = self._add_feedback(problem, check_result)

        trajectory.total_time = time.time() - start_time

        if self._verbose:
            print(f"\n{'='*50}")
            print(f"Result: {'SUCCESS' if trajectory.success else 'COMPLETED'}")
            print(f"Time: {trajectory.total_time:.2f}s")
            print(f"{'='*50}")

        return PSResult(
            success=trajectory.success,
            answer=solution.answer if solution else "",
            plan=plan,
            solution=solution,
            check_result=check_result,
            trajectory=trajectory,
        )

    def _add_feedback(self, problem: str, check_result: CheckResult) -> str:
        """Add feedback to problem for refinement.

        Args:
            problem: Original problem
            check_result: Check result with issues

        Returns:
            Problem with feedback
        """
        feedback = "\n\nFeedback from previous attempt:\n"
        for issue in check_result.issues:
            feedback += f"- {issue}\n"
        for suggestion in check_result.suggestions:
            feedback += f"- Suggestion: {suggestion}\n"
        return problem + feedback

    def quick_solve(self, problem: str) -> str:
        """Quick solve returning just the answer.

        Args:
            problem: Problem description

        Returns:
            Solution answer
        """
        result = self.solve(problem)
        return result.answer