from __future__ import annotations

from datetime import UTC, datetime
from pathlib import Path
from typing import Any

from pydantic import BaseModel, Field

from ai_orchestrator.schemas.task_request import TaskConstraints
from ai_orchestrator.schemas.task_status import TaskLifecycleStatus
from ai_orchestrator.schemas.workflow_models import (
    FocusIssueInventory,
    FocusReviewResult,
    MemoryHit,
    StructuredError,
    SubtaskRecord,
    ValidationSummary,
)


def utc_now() -> datetime:
    return datetime.now(UTC)


class ArtifactRecord(BaseModel):
    type: str
    path: str
    description: str | None = None


class CommandRecord(BaseModel):
    tool: str
    command: str
    success: bool
    stdout: str = ""
    stderr: str = ""
    exit_code: int | None = None
    structured_errors: list[StructuredError] = Field(default_factory=list)
    started_at: datetime = Field(default_factory=utc_now)
    completed_at: datetime = Field(default_factory=utc_now)


class StepRecord(BaseModel):
    step: str
    role: str
    status: str
    details: str | None = None
    started_at: datetime = Field(default_factory=utc_now)
    completed_at: datetime = Field(default_factory=utc_now)


class TestResult(BaseModel):
    command: str | None = None
    success: bool
    stdout: str = ""
    stderr: str = ""
    exit_code: int | None = None
    structured_errors: list[StructuredError] = Field(default_factory=list)
    started_at: datetime = Field(default_factory=utc_now)
    completed_at: datetime = Field(default_factory=utc_now)


class TaskTimestamps(BaseModel):
    created_at: datetime = Field(default_factory=utc_now)
    started_at: datetime | None = None
    updated_at: datetime = Field(default_factory=utc_now)
    completed_at: datetime | None = None


class TaskState(BaseModel):
    task_id: str
    goal: str
    repo_path: Path
    task_mode: str = "general_coding"
    workspace_path: Path | None = None
    workspace_strategy: str | None = None
    constraints: TaskConstraints
    attempt: int = 0
    current_step: str | None = None
    plan_steps: list[str] = Field(default_factory=list)
    repo_context: dict[str, Any] = Field(default_factory=dict)
    focus_issue_inventory: FocusIssueInventory | None = None
    focus_review_history: list[FocusReviewResult] = Field(default_factory=list)
    memory_hits: list[MemoryHit] = Field(default_factory=list)
    retry_focus_files: list[str] = Field(default_factory=list)
    rag_retrieval_events: list[dict[str, Any]] = Field(default_factory=list)
    validation_history: list[ValidationSummary] = Field(default_factory=list)
    subtasks: list[SubtaskRecord] = Field(default_factory=list)
    reflection_issues: list[str] = Field(default_factory=list)
    step_history: list[StepRecord] = Field(default_factory=list)
    command_history: list[CommandRecord] = Field(default_factory=list)
    test_results: list[TestResult] = Field(default_factory=list)
    artifacts: list[ArtifactRecord] = Field(default_factory=list)
    summary: str | None = None
    status: TaskLifecycleStatus = TaskLifecycleStatus.queued
    error_message: str | None = None
    timestamps: TaskTimestamps = Field(default_factory=TaskTimestamps)

    def execution_repo_path(self) -> Path:
        return self.workspace_path or self.repo_path

    def touch(self) -> None:
        self.timestamps.updated_at = utc_now()

    def mark_running(self, step: str | None = None) -> None:
        self.status = TaskLifecycleStatus.running
        self.current_step = step
        if self.timestamps.started_at is None:
            self.timestamps.started_at = utc_now()
        self.touch()

    def mark_terminal(
        self,
        status: TaskLifecycleStatus,
        *,
        summary: str | None = None,
        error_message: str | None = None,
    ) -> None:
        self.status = status
        if summary is not None:
            self.summary = summary
        if error_message is not None:
            self.error_message = error_message
        self.timestamps.completed_at = utc_now()
        self.touch()
