"""Solver for executing plans."""

from typing import Callable, Dict, List, Optional, Any

from .types import Plan, PlanStep, Solution, SolveStep, StepStatus


class Solver:
    """Executes plans step by step.

    The Solver takes a plan and executes each step,
    tracking intermediate results and building the solution.
    """

    def __init__(
        self,
        solve_fn: Optional[Callable[[str, Dict[str, Any]], str]] = None,
        verbose: bool = False,
    ):
        """Initialize the solver.

        Args:
            solve_fn: Optional custom solving function.
                     Takes (step_description, context) and returns result.
            verbose: Whether to print debug info
        """
        self._solve_fn = solve_fn or self._default_solve
        self._verbose = verbose

    def execute_plan(self, plan: Plan) -> Solution:
        """Execute a plan and produce a solution.

        Args:
            plan: The plan to execute

        Returns:
            A Solution object
        """
        solution = Solution(answer="", reasoning="")

        if self._verbose:
            print(f"\n[Solver] Executing plan with {len(plan.steps)} steps")

        # Context accumulates intermediate results
        context: Dict[str, Any] = {
            "problem": plan.problem,
            "step_results": {},
        }

        # Execute steps in order (assuming dependencies are met)
        for step in plan.steps:
            step.status = StepStatus.IN_PROGRESS

            if self._verbose:
                print(f"[Solver] Step {step.step_number}: {step.description[:50]}...")

            # Build context string
            context_str = self._build_context(context, step)

            # Execute step
            try:
                result = self._solve_fn(step.description, context)
                step.result = result
                step.status = StepStatus.COMPLETED
                context["step_results"][step.step_number] = result

                # Record solve step
                solve_step = SolveStep(
                    plan_step=step,
                    execution=result,
                    intermediate_results={"context": context_str},
                    success=True,
                )
                solution.steps.append(solve_step)

            except Exception as e:
                step.result = str(e)
                step.status = StepStatus.FAILED

                solve_step = SolveStep(
                    plan_step=step,
                    execution=str(e),
                    success=False,
                )
                solution.steps.append(solve_step)

                if self._verbose:
                    print(f"[Solver] Step failed: {e}")

                break

        # Build final answer from last completed step
        completed_steps = [s for s in solution.steps if s.success]
        if completed_steps:
            solution.answer = completed_steps[-1].execution
            solution.reasoning = "\n".join(
                f"Step {s.plan_step.step_number}: {s.execution[:100]}..."
                for s in completed_steps
            )

        if self._verbose:
            print(f"[Solver] Completed {len(completed_steps)}/{len(plan.steps)} steps")

        return solution

    def _build_context(self, context: Dict[str, Any], step: PlanStep) -> str:
        """Build context string for a step.

        Args:
            context: Execution context
            step: Current step

        Returns:
            Context string
        """
        parts = [f"Problem: {context['problem']}"]

        # Add previous step results
        for step_num, result in context.get("step_results", {}).items():
            parts.append(f"Step {step_num} result: {result[:200]}")

        parts.append(f"Current step: {step.description}")

        return "\n".join(parts)

    def _default_solve(self, step_description: str, context: Dict[str, Any]) -> str:
        """Default solving function.

        In real usage, this would call an LLM.
        """
        # Simulate solving
        problem = context.get("problem", "")
        return f"Executed: {step_description[:30]}... for problem: {problem[:30]}..."