"""
Core domain models — the runtime state of orchestration.

These are distinct from the *spec* models: the spec is the input contract,
these models track execution state as agents work.
"""

from __future__ import annotations

import hashlib
import json
from datetime import datetime
from enum import Enum

from pydantic import BaseModel, Field, computed_field


# ---------------------------------------------------------------------------
# Enums
# ---------------------------------------------------------------------------

class TaskStatus(str, Enum):
    PENDING = "pending"
    BLOCKED = "blocked"      # waiting on dependencies
    READY = "ready"          # all deps satisfied, waiting for agent
    IN_PROGRESS = "in_progress"
    VERIFYING = "verifying"  # task output is being verified
    PASSED = "passed"        # verification succeeded
    FAILED = "failed"        # verification failed (may retry)
    SKIPPED = "skipped"
    CANCELLED = "cancelled"


class WorkstreamStatus(str, Enum):
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class WorkflowStatus(str, Enum):
    PENDING = "pending"
    PLANNING = "planning"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


# ---------------------------------------------------------------------------
# Token tracking (WS-1.1)
# ---------------------------------------------------------------------------

class TokenUsage(BaseModel):
    """Detailed token usage from a single LLM generation."""
    input_tokens: int = 0
    output_tokens: int = 0
    cache_read_tokens: int = 0
    cache_creation_tokens: int = 0

    @computed_field  # type: ignore[prop-decorator]
    @property
    def total(self) -> int:
        return self.input_tokens + self.output_tokens

    def __add__(self, other: "TokenUsage") -> "TokenUsage":
        return TokenUsage(
            input_tokens=self.input_tokens + other.input_tokens,
            output_tokens=self.output_tokens + other.output_tokens,
            cache_read_tokens=self.cache_read_tokens + other.cache_read_tokens,
            cache_creation_tokens=self.cache_creation_tokens + other.cache_creation_tokens,
        )


# ---------------------------------------------------------------------------
# Verification feedback (WS-1.2)
# ---------------------------------------------------------------------------

class TestCaseResult(BaseModel):
    """Result of a single test case from structured test parsing."""
    __test__ = False  # Prevent pytest collection
    name: str
    passed: bool
    error_message: str = ""


class VerificationFeedback(BaseModel):
    """Feedback from a single verification attempt, stored for retry context."""
    attempt: int
    verifier_type: str  # "llm_judge", "test_runner", "chained"
    passed: bool
    score: float = 0.0
    feedback: str = ""
    suggested_improvements: list[str] = Field(default_factory=list)
    test_results: list[TestCaseResult] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Task
# ---------------------------------------------------------------------------

class ArtifactManifestEntry(BaseModel):
    """Metadata for a file produced by a task — path + size, no content."""
    path: str
    size_bytes: int
    file_type: str = ""  # extension: "py", "md", "json", etc.


class TaskResult(BaseModel):
    """Output produced by an agent for a single task."""
    output: str = ""
    artifacts: dict[str, str] = Field(
        default_factory=dict,
        description="name → content mapping of files/artifacts produced",
    )
    file_manifest: list[ArtifactManifestEntry] = Field(
        default_factory=list,
        description="Authoritative record of files produced (path + size, no content)",
    )
    generated_tests: list[GeneratedTest] = Field(default_factory=list)
    error: str | None = None
    token_usage: int = 0
    token_usage_detailed: TokenUsage | None = None
    wall_seconds: float = 0.0
    learnings: list[str] = Field(default_factory=list)


class GeneratedTest(BaseModel):
    """A test generated by an agent to verify its own output."""
    filename: str
    content: str
    test_type: str = "unit"
    framework: str = "pytest"


class StructuredTaskPrompt(BaseModel):
    """Structured format for task descriptions, validated by the plan checker."""
    objective: str = Field(default="", description="What must be accomplished")
    files: list[str] = Field(default_factory=list, description="File scope — only touch these")
    action: str = Field(default="", description="Specific steps to take")
    verify: str = Field(default="", description="How to verify the output is correct")
    done: str = Field(default="", description="Definition of done")


class Task(BaseModel):
    """Smallest unit of work — assigned to a single agent."""
    id: str
    workstream_id: str
    workflow_id: str

    title: str
    description: str
    structured_prompt: StructuredTaskPrompt | None = None
    status: TaskStatus = TaskStatus.PENDING

    assigned_agent_id: str | None = None
    depends_on: list[str] = Field(
        default_factory=list,
        description="Task IDs that must complete before this task starts",
    )
    acceptance_criteria_ids: list[str] = Field(
        default_factory=list,
        description="Which spec-level acceptance criteria this task contributes to",
    )

    # Verification
    verification_strategy: str = Field(
        default="llm_judge",
        description="'llm_judge', 'test', 'manual', or custom verifier name",
    )
    skeleton_tests: list[str] = Field(
        default_factory=list,
        description="Skeleton test code the agent must implement and pass",
    )

    # Execution tracking
    result: TaskResult | None = None
    attempt: int = 0
    max_retries: int = 3
    created_at: datetime = Field(default_factory=datetime.utcnow)
    started_at: datetime | None = None
    completed_at: datetime | None = None

    # Verification feedback for retry context (WS-1.2)
    attempt_feedback: list[VerificationFeedback] = Field(default_factory=list)

    # Linear sync (WS-1.3)
    linear_issue_id: str | None = None

    # Transient — set by orchestrator before dispatch, not persisted
    enriched_prompt: str | None = Field(
        default=None,
        exclude=True,
        description="Pre-built context prompt from ContextBuilder. Agents use this when set.",
    )

    def content_hash(self) -> str:
        """Hash for dedup — two tasks with the same hash are duplicates."""
        payload = json.dumps(
            {"title": self.title, "description": self.description, "workstream_id": self.workstream_id},
            sort_keys=True,
        )
        return hashlib.sha256(payload.encode()).hexdigest()[:16]

    @property
    def is_terminal(self) -> bool:
        return self.status in (
            TaskStatus.PASSED,
            TaskStatus.FAILED,
            TaskStatus.SKIPPED,
            TaskStatus.CANCELLED,
        )


# ---------------------------------------------------------------------------
# Workstream
# ---------------------------------------------------------------------------

class Workstream(BaseModel):
    """A logical grouping of related tasks (e.g. 'Backend API', 'Frontend UI')."""
    id: str
    workflow_id: str
    name: str
    description: str
    status: WorkstreamStatus = WorkstreamStatus.PENDING
    task_ids: list[str] = Field(default_factory=list)
    created_at: datetime = Field(default_factory=datetime.utcnow)


# ---------------------------------------------------------------------------
# Workflow
# ---------------------------------------------------------------------------

class Workflow(BaseModel):
    """
    A complete execution plan derived from a Specification.
    """
    id: str
    spec_id: str
    status: WorkflowStatus = WorkflowStatus.PENDING
    workstream_ids: list[str] = Field(default_factory=list)

    spec_content_hash: str = Field(
        default="",
        description="Hash of the spec that produced this workflow, for cache/dedup",
    )
    created_at: datetime = Field(default_factory=datetime.utcnow)
    completed_at: datetime | None = None

    # Progress tracking
    total_tasks: int = 0
    completed_tasks: int = 0
    failed_tasks: int = 0
    replan_count: int = 0

    # Linear sync (WS-1.3)
    linear_project_id: str | None = None


# ---------------------------------------------------------------------------
# WorkflowState — the aggregate root
# ---------------------------------------------------------------------------

class WorkflowState(BaseModel):
    """
    Single source of truth for all orchestration state.

    This is what gets persisted to the state backend and is the only object
    the orchestrator reads/writes during execution.
    """
    workflows: dict[str, Workflow] = Field(default_factory=dict)
    workstreams: dict[str, Workstream] = Field(default_factory=dict)
    tasks: dict[str, Task] = Field(default_factory=dict)

    # Dedup registry: content_hash → task_id
    task_hashes: dict[str, str] = Field(default_factory=dict)

    def get_ready_tasks(self, workflow_id: str) -> list[Task]:
        """Return tasks whose dependencies are all satisfied."""
        ready = []
        for task in self.tasks.values():
            if task.workflow_id != workflow_id:
                continue
            if task.status != TaskStatus.PENDING:
                continue
            deps_satisfied = all(
                self.tasks[dep_id].status == TaskStatus.PASSED
                for dep_id in task.depends_on
                if dep_id in self.tasks
            )
            if deps_satisfied:
                ready.append(task)
        return ready

    def is_workflow_complete(self, workflow_id: str) -> bool:
        wf = self.workflows.get(workflow_id)
        if not wf:
            return False
        tasks = [t for t in self.tasks.values() if t.workflow_id == workflow_id]
        if not tasks:
            return False
        return all(t.is_terminal for t in tasks)

    def is_workflow_failed(self, workflow_id: str) -> bool:
        tasks = [t for t in self.tasks.values() if t.workflow_id == workflow_id]
        return any(t.status == TaskStatus.FAILED for t in tasks)

    def is_workflow_stalled(self, workflow_id: str) -> bool:
        """True when no more progress is possible.

        A workflow is stalled when no pending task can become ready — every
        pending task is transitively blocked by at least one FAILED task.
        """
        pending = [
            t for t in self.tasks.values()
            if t.workflow_id == workflow_id and t.status == TaskStatus.PENDING
        ]
        if not pending:
            return self.is_workflow_complete(workflow_id)

        # A task is "blocked" if any of its dependencies are FAILED or blocked.
        # Use memoized check to avoid recomputation.
        blocked_cache: dict[str, bool] = {}

        def _is_blocked(task_id: str) -> bool:
            if task_id in blocked_cache:
                return blocked_cache[task_id]
            task = self.tasks.get(task_id)
            if not task:
                return False
            if task.status == TaskStatus.FAILED:
                blocked_cache[task_id] = True
                return True
            if task.status == TaskStatus.PASSED:
                blocked_cache[task_id] = False
                return False
            # PENDING: blocked if any dependency is blocked
            result = any(_is_blocked(dep_id) for dep_id in task.depends_on)
            blocked_cache[task_id] = result
            return result

        for task in pending:
            if not task.depends_on:
                return False  # No deps → can run → not stalled
            if not _is_blocked(task.id):
                return False  # Not blocked → can eventually run

        return True  # All pending tasks are transitively blocked by failures

    def register_task(self, task: Task) -> Task | None:
        """Register a task, deduplicating by content hash. Returns None if duplicate."""
        h = task.content_hash()
        if h in self.task_hashes:
            return None
        self.task_hashes[h] = task.id
        self.tasks[task.id] = task
        return task
