"""Shared data types, enums, and abstract base classes for the Pointer multi-agent system."""

from __future__ import annotations

import uuid
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum


class Phase(str, Enum):
    FEASIBILITY_CHECK = "feasibility_check"
    PLANNING = "planning"
    EXECUTING = "executing"
    GUI = "gui"
    VERIFYING = "verifying"
    REPLANNING = "replanning"
    DONE = "done"
    FAILED = "failed"


class ExecutorDecision(str, Enum):
    DONE = "done"
    FAIL = "fail"
    GUI = "gui"


@dataclass
class ExecutorAction:
    decision: ExecutorDecision
    reason: str


class MilestoneStatus(str, Enum):
    PENDING = "pending"
    COMPLETED = "completed"


class VerifierDecision(str, Enum):
    DONE = "done"
    FAILED = "failed"
    INFEASIBLE = "infeasible"


@dataclass
class SubMilestone:
    description: str
    done: bool = False

    def to_dict(self) -> dict:
        return {"description": self.description, "done": self.done}

    @classmethod
    def from_dict(cls, d: dict) -> SubMilestone:
        return cls(description=d["description"], done=d.get("done", False))


@dataclass
class Milestone:
    description: str
    status: MilestoneStatus = MilestoneStatus.PENDING
    id: str = field(default_factory=lambda: uuid.uuid4().hex[:8])
    sub_milestones: list[SubMilestone] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "description": self.description,
            "status": self.status.value,
            "sub_milestones": [sm.to_dict() for sm in self.sub_milestones],
        }

    @classmethod
    def from_dict(cls, d: dict) -> Milestone:
        return cls(
            id=d.get("id", uuid.uuid4().hex[:8]),
            description=d["description"],
            status=MilestoneStatus(d.get("status", "pending")),
            sub_milestones=[SubMilestone.from_dict(sm) for sm in d.get("sub_milestones", [])],
        )


@dataclass
class Plan:
    task_instruction: str
    milestones: list[Milestone] = field(default_factory=list)
    current_milestone_idx: int = 0
    revision_count: int = 0
    infeasible_reason: str | None = None

    def current_milestone(self) -> Milestone:
        if self.current_milestone_idx >= len(self.milestones):
            return Milestone("All milesteones completed!", status=MilestoneStatus.COMPLETED)
        return self.milestones[self.current_milestone_idx]

    def mark_current_completed(self) -> Milestone | None:
        m = self.current_milestone()
        m.status = MilestoneStatus.COMPLETED
        self.current_milestone_idx += 1
        return m

    def update_milestone(self, idx_1based: int, new_description: str) -> str | None:
        idx = idx_1based - 1
        if idx < 0 or idx >= len(self.milestones):
            return f"Invalid milestone index {idx_1based}. Valid range: 1-{len(self.milestones)}."
        m = self.milestones[idx]
        if m.status == MilestoneStatus.COMPLETED:
            return f"Cannot update milestone {idx_1based}: it is already completed."
        m.description = new_description
        return None

    def add_milestone(self, description: str, idx_1based: int | None = None) -> str | None:
        new_m = Milestone(description=description)
        if idx_1based is None:
            self.milestones.append(new_m)
            return None
        idx = idx_1based - 1
        if idx < 0 or idx > len(self.milestones):
            return f"Invalid index {idx_1based}. Valid range: 1-{len(self.milestones) + 1}."
        if idx < self.current_milestone_idx:
            return f"Cannot insert at position {idx_1based}: positions 1-{self.current_milestone_idx} are already completed."
        self.milestones.insert(idx, new_m)
        if idx <= self.current_milestone_idx:
            self.current_milestone_idx += 1
        return None

    def delete_milestone(self, idx_1based: int) -> str | None:
        idx = idx_1based - 1
        if idx < 0 or idx >= len(self.milestones):
            return f"Invalid milestone index {idx_1based}. Valid range: 1-{len(self.milestones)}."
        if len(self.milestones) <= 1:
            return "Cannot delete the last remaining milestone."
        m = self.milestones[idx]
        if m.status == MilestoneStatus.COMPLETED:
            return f"Cannot delete milestone {idx_1based}: it is already completed."
        if idx == self.current_milestone_idx:
            return "Cannot delete the current milestone — update it or mark it done instead."
        self.milestones.pop(idx)
        if idx < self.current_milestone_idx:
            self.current_milestone_idx -= 1
        return None

    def set_sub_milestones(self, sub_milestones: list[str]) -> str | None:
        if self.current_milestone_idx >= len(self.milestones):
            return "All milestones completed — cannot set sub-milestones."
        m = self.milestones[self.current_milestone_idx]
        m.sub_milestones = [SubMilestone(description=desc) for desc in sub_milestones]
        return None

    def mark_sub_milestone_done(self, sub_idx_1based: int) -> str | None:
        if self.current_milestone_idx >= len(self.milestones):
            return "All milestones completed."
        m = self.milestones[self.current_milestone_idx]
        idx = sub_idx_1based - 1
        if idx < 0 or idx >= len(m.sub_milestones):
            return f"Invalid sub-milestone index {sub_idx_1based}. Valid range: 1-{len(m.sub_milestones)}."
        m.sub_milestones[idx].done = True
        return None

    def format(self) -> str:
        if not self.milestones:
            return "(no plan)"
        lines: list[str] = []
        for i, m in enumerate(self.milestones):
            if m.status == MilestoneStatus.COMPLETED:
                marker = "[X]"
            elif i == self.current_milestone_idx:
                marker = "[>]"
            else:
                marker = "[ ]"
            lines.append(f"{marker} {i + 1}. {m.description}")
            if i == self.current_milestone_idx and m.sub_milestones:
                for j, sm in enumerate(m.sub_milestones):
                    sm_marker = "[x]" if sm.done else "[ ]"
                    lines.append(f"     {sm_marker} {j + 1}. {sm.description}")

        if self.infeasible_reason:
            lines.append(
                f"\nWARNING: The planner suggests this might be infeasible."
                f"\nREASON: {self.infeasible_reason}"
                f"\nHOWEVER: Try to work through the milestones and declare the"
                " task infeasible if the task does happen to be infeasible."
            )

        return "\n".join(lines)

    def to_dict(self) -> dict:
        return {
            "task_instruction": self.task_instruction,
            "milestones": [m.to_dict() for m in self.milestones],
            "current_milestone_idx": self.current_milestone_idx,
            "revision_count": self.revision_count,
            "infeasible_reason": self.infeasible_reason,
        }

    @classmethod
    def from_dict(cls, d: dict) -> Plan:
        return cls(
            task_instruction=d["task_instruction"],
            milestones=[Milestone.from_dict(m) for m in d.get("milestones", [])],
            current_milestone_idx=d.get("current_milestone_idx", 0),
            revision_count=d.get("revision_count", 0),
            infeasible_reason=d.get("infeasible_reason"),
        )


@dataclass
class VerificationResult:
    decision: VerifierDecision
    reasoning: str = ""
    diagnosis: str = ""
    suggestion: str = ""

    def __str__(self) -> str:
        s = "Verification Result:\n"
        s += f"Decision: {self.decision}\n"
        s += f"Reasoning: {self.reasoning}\n"
        s += f"Diagnosis: {self.diagnosis}\n"
        s += f"Suggestion: {self.suggestion}\n"
        return s


@dataclass
class ActionRecord:
    index: int
    tool_name: str
    tool_input: str
    tool_output: str
    timestamp: str
    thinking: str = ""
    reasoning: str = ""
    milestone_id: str = ""


@dataclass
class FeasibilityResult:
    infeasible: bool
    reason: str


@dataclass
class AgentState:
    phase: Phase = Phase.FEASIBILITY_CHECK
    step_count: int = 0
    plan: Plan | None = None
    action_records: list[ActionRecord] = field(default_factory=list)
    latest_screenshot_b64: str = ""
    replan_guidance: str = ""
    executor_action: ExecutorAction | None = None
    verification_result: VerificationResult | None = None

    def log_action(
        self,
        tool_name: str,
        tool_input: str,
        tool_output: str,
        thinking: str = "",
        reasoning: str = "",
        milestone_id: str = "",
    ) -> None:
        record = ActionRecord(
            index=len(self.action_records) + 1,
            tool_name=tool_name,
            tool_input=tool_input,
            tool_output=tool_output,
            timestamp=datetime.now().isoformat(),
            thinking=thinking,
            reasoning=reasoning,
            milestone_id=milestone_id,
        )
        self.action_records.append(record)

    def recent_log(self, n: int = 5) -> str:
        records = self.action_records[-n:]
        if not records:
            return "(no actions taken yet)"
        lines = []
        for r in records:
            lines.append(f"  #{r.index} [{r.tool_name}]: {r.tool_input[:100]}")
            output_preview = r.tool_output[:200] + "..." if len(r.tool_output) > 200 else r.tool_output
            lines.append(f"    Result: {output_preview}")
        return "\n".join(lines)
