"""Core types for Plan-and-Solve framework."""

from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Dict, List, Optional
import uuid


class StepStatus(Enum):
    """Status of a planning/solving step."""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class PlanStep:
    """A single step in a plan."""
    description: str
    step_number: int
    status: StepStatus = StepStatus.PENDING
    result: str = ""
    dependencies: List[int] = field(default_factory=list)
    step_id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])

    def to_dict(self) -> Dict[str, Any]:
        return {
            "step_id": self.step_id,
            "description": self.description,
            "step_number": self.step_number,
            "status": self.status.value,
            "result": self.result,
            "dependencies": self.dependencies,
        }


@dataclass
class Plan:
    """A complete plan for solving a problem."""
    problem: str
    steps: List[PlanStep] = field(default_factory=list)
    plan_id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])
    created_at: datetime = field(default_factory=datetime.now)

    def add_step(self, description: str, dependencies: Optional[List[int]] = None) -> PlanStep:
        """Add a step to the plan.

        Args:
            description: Step description
            dependencies: List of step numbers this depends on

        Returns:
            The created step
        """
        step = PlanStep(
            description=description,
            step_number=len(self.steps) + 1,
            dependencies=dependencies or [],
        )
        self.steps.append(step)
        return step

    def get_step(self, step_number: int) -> Optional[PlanStep]:
        """Get a step by number.

        Args:
            step_number: Step number (1-indexed)

        Returns:
            The step if found
        """
        if 1 <= step_number <= len(self.steps):
            return self.steps[step_number - 1]
        return None

    def get_next_pending(self) -> Optional[PlanStep]:
        """Get the next pending step whose dependencies are met.

        Returns:
            Next executable step
        """
        for step in self.steps:
            if step.status != StepStatus.PENDING:
                continue

            # Check dependencies
            deps_met = all(
                self.get_step(d).status == StepStatus.COMPLETED
                for d in step.dependencies
                if self.get_step(d) is not None
            )

            if deps_met:
                return step

        return None

    def is_complete(self) -> bool:
        """Check if all steps are completed.

        Returns:
            True if complete
        """
        return all(s.status == StepStatus.COMPLETED for s in self.steps)

    def to_dict(self) -> Dict[str, Any]:
        return {
            "plan_id": self.plan_id,
            "problem": self.problem,
            "steps": [s.to_dict() for s in self.steps],
            "created_at": self.created_at.isoformat(),
        }

    def to_prompt(self) -> str:
        """Format plan for use in prompts.

        Returns:
            Formatted plan string
        """
        lines = [f"Plan for: {self.problem}\n"]
        for step in self.steps:
            status_icon = {
                StepStatus.PENDING: "○",
                StepStatus.IN_PROGRESS: "◐",
                StepStatus.COMPLETED: "●",
                StepStatus.FAILED: "✗",
                StepStatus.SKIPPED: "⊘",
            }.get(step.status, "○")

            lines.append(f"  {status_icon} Step {step.step_number}: {step.description}")

        return "\n".join(lines)


@dataclass
class SolveStep:
    """Result of executing a plan step."""
    plan_step: PlanStep
    execution: str
    intermediate_results: Dict[str, Any] = field(default_factory=dict)
    success: bool = True

    def to_dict(self) -> Dict[str, Any]:
        return {
            "step_number": self.plan_step.step_number,
            "description": self.plan_step.description,
            "execution": self.execution,
            "success": self.success,
            "intermediate_results": self.intermediate_results,
        }


@dataclass
class Solution:
    """Complete solution to a problem."""
    answer: str
    steps: List[SolveStep] = field(default_factory=list)
    reasoning: str = ""
    confidence: float = 1.0
    solution_id: str = field(default_factory=lambda: str(uuid.uuid4())[:8])

    def to_dict(self) -> Dict[str, Any]:
        return {
            "solution_id": self.solution_id,
            "answer": self.answer,
            "reasoning": self.reasoning,
            "confidence": self.confidence,
            "steps": [s.to_dict() for s in self.steps],
        }


@dataclass
class CheckResult:
    """Result of self-checking a solution."""
    passed: bool
    issues: List[str] = field(default_factory=list)
    suggestions: List[str] = field(default_factory=list)
    score: float = 1.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "passed": self.passed,
            "issues": self.issues,
            "suggestions": self.suggestions,
            "score": self.score,
        }


@dataclass
class PSTrajectory:
    """Complete trajectory of a Plan-and-Solve session."""
    problem: str
    plan: Optional[Plan] = None
    solution: Optional[Solution] = None
    check_result: Optional[CheckResult] = None
    success: bool = False
    total_time: float = 0.0

    def to_dict(self) -> Dict[str, Any]:
        return {
            "problem": self.problem,
            "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,
            "success": self.success,
            "total_time": self.total_time,
        }

    def to_markdown(self) -> str:
        """Export as markdown."""
        lines = [
            "# Plan-and-Solve Trajectory\n",
            f"## Problem\n{self.problem}\n",
        ]

        if self.plan:
            lines.append("## Plan\n")
            for step in self.plan.steps:
                icon = "✓" if step.status == StepStatus.COMPLETED else "○"
                lines.append(f"{icon} **Step {step.step_number}**: {step.description}")
                if step.result:
                    lines.append(f"   Result: {step.result[:100]}...")
            lines.append("")

        if self.solution:
            lines.extend([
                "## Solution\n",
                f"**Answer**: {self.solution.answer}\n",
                "**Reasoning**:",
                self.solution.reasoning or "N/A",
                "",
            ])

        if self.check_result:
            lines.extend([
                "## Self-Check\n",
                f"**Passed**: {self.check_result.passed}",
                f"**Score**: {self.check_result.score:.2f}",
                "",
            ])

        return "\n".join(lines)