"""Core data models for Dobby meta-agent orchestrator."""

from __future__ import annotations

import secrets
from dataclasses import dataclass, field
from datetime import UTC, datetime
from enum import Enum
from typing import Any

# Alias to avoid shadowing by dataclass fields named "timezone"
timezone_utc = UTC


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

class TodoStatus(Enum):
    OPEN = "open"
    DONE = "done"
    CANCELLED = "cancelled"


class TaskStatus(Enum):
    PENDING = "pending"
    DRAFTING = "drafting"
    ACTIVE = "active"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"
    PAUSED = "paused"


# Legal state transitions: from_status -> {allowed target statuses}
_VALID_TRANSITIONS: dict[TaskStatus, set[TaskStatus]] = {
    TaskStatus.PENDING: {TaskStatus.RUNNING, TaskStatus.CANCELLED},
    TaskStatus.DRAFTING: {TaskStatus.PENDING, TaskStatus.RUNNING, TaskStatus.CANCELLED},
    TaskStatus.ACTIVE: {
        TaskStatus.RUNNING, TaskStatus.PENDING, TaskStatus.CANCELLED, TaskStatus.PAUSED
    },
    TaskStatus.RUNNING: {
        TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.CANCELLED,
        TaskStatus.ACTIVE, TaskStatus.PAUSED,
    },
    TaskStatus.COMPLETED: {TaskStatus.PENDING, TaskStatus.ACTIVE},  # via retry
    TaskStatus.FAILED: {TaskStatus.PENDING, TaskStatus.ACTIVE},     # via retry
    TaskStatus.CANCELLED: {TaskStatus.PENDING, TaskStatus.ACTIVE},  # via retry
    TaskStatus.PAUSED: {TaskStatus.ACTIVE, TaskStatus.CANCELLED},   # resume or cancel
}


class InvalidTransitionError(Exception):
    """Raised when an illegal status transition is attempted."""


def validate_transition(current: TaskStatus, target: TaskStatus) -> None:
    """Enforce legal status transitions.

    Raises InvalidTransitionError if the transition is not allowed.
    """
    if current == target:
        return  # no-op transitions are always fine
    allowed = _VALID_TRANSITIONS.get(current, set())
    if target not in allowed:
        raise InvalidTransitionError(
            f"Cannot transition from {current.value!r} to {target.value!r}"
        )


class TaskMode(Enum):
    QUICK = "quick"
    PLANNED = "planned"


class Priority(Enum):
    LOW = "low"
    NORMAL = "normal"
    HIGH = "high"
    URGENT = "urgent"

    @property
    def weight(self) -> int:
        _weights = {
            Priority.LOW: 10,
            Priority.NORMAL: 20,
            Priority.HIGH: 30,
            Priority.URGENT: 40,
        }
        return _weights[self]


class ScheduleType(Enum):
    ONCE = "once"
    CRON = "cron"
    INTERVAL = "interval"
    AFTER = "after"


# ---------------------------------------------------------------------------
# Schedule
# ---------------------------------------------------------------------------

@dataclass
class Schedule:
    type: ScheduleType = ScheduleType.ONCE
    expr: str | None = None        # cron expression
    every: str | None = None       # interval string e.g. "30m", "2h"
    until: str | None = None       # ISO datetime string
    depends_on: str | None = None  # task id for AFTER type

    def interval_seconds(self) -> int:
        """Parse an interval string like '30m' or '2h' into seconds."""
        if self.every is None:
            raise ValueError("interval_seconds() requires 'every' to be set")
        s = self.every.strip()
        if s.endswith("m"):
            return int(s[:-1]) * 60
        if s.endswith("h"):
            return int(s[:-1]) * 3600
        if s.endswith("s"):
            return int(s[:-1])
        if s.endswith("d"):
            return int(s[:-1]) * 86400
        raise ValueError(f"Cannot parse interval: {self.every!r}")

    def to_dict(self) -> dict[str, Any]:
        return {
            "type": self.type.value,
            "expr": self.expr,
            "every": self.every,
            "until": self.until,
            "depends_on": self.depends_on,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Schedule:
        return cls(
            type=ScheduleType(d["type"]),
            expr=d.get("expr"),
            every=d.get("every"),
            until=d.get("until"),
            depends_on=d.get("depends_on"),
        )


# ---------------------------------------------------------------------------
# Artifact
# ---------------------------------------------------------------------------

@dataclass
class Artifact:
    """A structured output artifact from a task execution."""
    type: str          # "pr", "log", "diff", "file", "url", etc.
    path: str | None = None   # local file path
    url: str | None = None    # remote URL (e.g. PR link)
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"type": self.type}
        if self.path is not None:
            d["path"] = self.path
        if self.url is not None:
            d["url"] = self.url
        if self.metadata:
            d["metadata"] = self.metadata
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Artifact:
        return cls(
            type=d["type"],
            path=d.get("path"),
            url=d.get("url"),
            metadata=d.get("metadata", {}),
        )


# ---------------------------------------------------------------------------
# InfraConfig
# ---------------------------------------------------------------------------

@dataclass
class InfraConfig:
    name: str
    type: str                  # "ssh", "modal", "local", etc.
    host: str | None = None    # SSH target
    gpus: int | None = None    # number of GPUs
    conda_env: str | None = None
    gpu: str | None = None     # GPU model e.g. "H100"

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "type": self.type,
            "host": self.host,
            "gpus": self.gpus,
            "conda_env": self.conda_env,
            "gpu": self.gpu,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> InfraConfig:
        return cls(
            name=d["name"],
            type=d["type"],
            host=d.get("host"),
            gpus=d.get("gpus"),
            conda_env=d.get("conda_env"),
            gpu=d.get("gpu"),
        )


# ---------------------------------------------------------------------------
# ProjectPolicies
# ---------------------------------------------------------------------------

@dataclass
class ProjectPolicies:
    """Default policies for a project — merged into tasks at creation time."""
    default_backend: str | None = None       # "claude" or "codex"
    default_timeout: int | None = None       # seconds
    max_concurrent: int | None = None        # override global per-project lock
    default_priority: str | None = None      # "low", "normal", "high", "urgent"
    auto_chain: bool = False                 # auto-depend on latest pending/running task

    def to_dict(self) -> dict[str, Any]:
        return {
            "default_backend": self.default_backend,
            "default_timeout": self.default_timeout,
            "max_concurrent": self.max_concurrent,
            "default_priority": self.default_priority,
            "auto_chain": self.auto_chain,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> ProjectPolicies:
        return cls(
            default_backend=d.get("default_backend"),
            default_timeout=d.get("default_timeout"),
            max_concurrent=d.get("max_concurrent"),
            default_priority=d.get("default_priority"),
            auto_chain=d.get("auto_chain", False),
        )


# ---------------------------------------------------------------------------
# Project
# ---------------------------------------------------------------------------

@dataclass
class Project:
    name: str
    repo: str
    git_remote: str | None = None
    context: str | None = None
    claude_md: bool = True
    dispatch_to: str = "local"  # "local" or "remote:<hostname>"
    infra: list[InfraConfig] = field(default_factory=list)
    policies: ProjectPolicies = field(default_factory=ProjectPolicies)

    @property
    def is_remote(self) -> bool:
        return self.dispatch_to.startswith("remote:")

    @property
    def remote_host(self) -> str | None:
        if self.is_remote:
            return self.dispatch_to.split(":", 1)[1]
        return None

    def to_dict(self) -> dict[str, Any]:
        return {
            "name": self.name,
            "repo": self.repo,
            "git_remote": self.git_remote,
            "context": self.context,
            "claude_md": self.claude_md,
            "dispatch_to": self.dispatch_to,
            "infra": [i.to_dict() for i in self.infra],
            "policies": self.policies.to_dict(),
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Project:
        policies_raw = d.get("policies")
        policies = ProjectPolicies.from_dict(policies_raw) if policies_raw else ProjectPolicies()
        return cls(
            name=d["name"],
            repo=d["repo"],
            git_remote=d.get("git_remote"),
            context=d.get("context"),
            claude_md=d.get("claude_md", True),
            dispatch_to=d.get("dispatch_to", "local"),
            infra=[InfraConfig.from_dict(i) for i in d.get("infra", [])],
            policies=policies,
        )


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

def _generate_task_id() -> str:
    date_str = datetime.now(UTC).strftime("%Y%m%d")
    suffix = secrets.token_hex(2)  # 4-char hex
    return f"t-{date_str}-{suffix}"


def generate_batch_id() -> str:
    """Generate a unique batch ID like 'b-20260322-a1b2'."""
    date_str = datetime.now(UTC).strftime("%Y%m%d")
    suffix = secrets.token_hex(2)
    return f"b-{date_str}-{suffix}"


def _default_schedule() -> Schedule:
    return Schedule(type=ScheduleType.ONCE)


def _is_recurring(schedule: Schedule) -> bool:
    return schedule.type in (ScheduleType.CRON, ScheduleType.INTERVAL)


def _is_waiting(schedule: Schedule) -> bool:
    """AFTER tasks start as ACTIVE (waiting on dependency), not PENDING."""
    return schedule.type == ScheduleType.AFTER


@dataclass
class TaskEvent:
    """A single entry in the task's event journal."""
    status: str
    timestamp: str  # ISO format
    detail: str | None = None

    def to_dict(self) -> dict[str, Any]:
        d: dict[str, Any] = {"status": self.status, "timestamp": self.timestamp}
        if self.detail is not None:
            d["detail"] = self.detail
        return d

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> TaskEvent:
        return cls(
            status=d["status"],
            timestamp=d["timestamp"],
            detail=d.get("detail"),
        )


@dataclass
class Task:
    project: str
    description: str
    title: str = ""  # short concise title for sidebar display
    id: str = field(default_factory=_generate_task_id)
    mode: TaskMode = TaskMode.QUICK
    priority: Priority = Priority.NORMAL
    exit_criteria: str | None = None
    schedule: Schedule = field(default_factory=_default_schedule)
    infra: str | None = None
    timeout: int | None = None
    checkpoint: bool = False
    plan: list[str] = field(default_factory=list)
    session_id: str | None = None
    backend: str = "claude"  # "claude" or "codex"
    target: str = "local"  # spoke id: "local", "gpu-a100", etc.
    batch_id: str | None = None  # links task to a batch dispatch
    trace_id: str | None = None  # links task to a trajectory trace
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    started_at: datetime | None = None
    last_run_at: datetime | None = None
    next_run_at: datetime | None = None
    attempt_number: int = 0
    max_failures: int = 3
    consecutive_failures: int = 0
    events: list[TaskEvent] = field(default_factory=list)
    status: TaskStatus = field(init=False)

    def __post_init__(self) -> None:
        # Status is derived from mode and schedule — not set directly at init
        if _is_recurring(self.schedule):
            self.status = TaskStatus.ACTIVE
        elif _is_waiting(self.schedule):
            # AFTER tasks wait for their dependency — start as ACTIVE until resolved
            self.status = TaskStatus.ACTIVE
        elif self.mode == TaskMode.PLANNED:
            self.status = TaskStatus.DRAFTING
        else:
            self.status = TaskStatus.PENDING
        # Record the initial "created" event if the events list is empty
        if not self.events:
            self.events.append(TaskEvent(
                status=self.status.value,
                timestamp=self.created_at.isoformat(),
                detail="created",
            ))

    def record_event(self, status: TaskStatus, detail: str | None = None) -> None:
        """Append a state-transition event to the journal."""
        self.events.append(TaskEvent(
            status=status.value,
            timestamp=datetime.now(UTC).isoformat(),
            detail=detail,
        ))

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "project": self.project,
            "title": self.title,
            "description": self.description,
            "mode": self.mode.value,
            "priority": self.priority.value,
            "exit_criteria": self.exit_criteria,
            "schedule": self.schedule.to_dict(),
            "infra": self.infra,
            "timeout": self.timeout,
            "checkpoint": self.checkpoint,
            "plan": self.plan,
            "session_id": self.session_id,
            "backend": self.backend,
            "target": self.target,
            "batch_id": self.batch_id,
            "trace_id": self.trace_id,
            "created_at": self.created_at.isoformat(),
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "last_run_at": self.last_run_at.isoformat() if self.last_run_at else None,
            "next_run_at": self.next_run_at.isoformat() if self.next_run_at else None,
            "attempt_number": self.attempt_number,
            "max_failures": self.max_failures,
            "consecutive_failures": self.consecutive_failures,
            "events": [e.to_dict() for e in self.events],
            "status": self.status.value,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Task:
        events = [TaskEvent.from_dict(e) for e in d.get("events", [])]
        task = cls(
            id=d["id"],
            project=d["project"],
            description=d["description"],
            title=d.get("title", ""),
            mode=TaskMode(d["mode"]),
            priority=Priority(d["priority"]),
            exit_criteria=d.get("exit_criteria"),
            schedule=Schedule.from_dict(d["schedule"]),
            infra=d.get("infra"),
            timeout=d.get("timeout"),
            checkpoint=d.get("checkpoint", False),
            plan=d.get("plan", []),
            session_id=d.get("session_id"),
            backend=d.get("backend", "claude"),
            target=d.get("target", "local"),
            batch_id=d.get("batch_id"),
            trace_id=d.get("trace_id"),
            created_at=datetime.fromisoformat(d["created_at"]),
            attempt_number=d.get("attempt_number", 0),
            max_failures=d.get("max_failures", 3),
            consecutive_failures=d.get("consecutive_failures", 0),
            events=events,
        )
        started_at = d.get("started_at")
        last_run_at = d.get("last_run_at")
        next_run_at = d.get("next_run_at")
        task.started_at = datetime.fromisoformat(started_at) if started_at else None
        task.last_run_at = datetime.fromisoformat(last_run_at) if last_run_at else None
        task.next_run_at = datetime.fromisoformat(next_run_at) if next_run_at else None
        # Restore persisted status (overrides the derived default)
        task.status = TaskStatus(d["status"])
        return task


# ---------------------------------------------------------------------------
# Result
# ---------------------------------------------------------------------------

def _default_cost() -> dict[str, Any]:
    return {
        "input_tokens": 0,
        "output_tokens": 0,
        "total_tokens": 0,
        "cost_usd": 0.0,
    }


@dataclass
class Result:
    task_id: str
    status: str
    summary: str
    artifacts: list[Artifact] = field(default_factory=list)
    metrics: dict[str, Any] = field(default_factory=dict)
    gpu_info: dict[str, Any] | None = None
    cost: dict[str, Any] = field(default_factory=_default_cost)
    session_id: str | None = None  # Claude Code session ID for --resume
    started_at: datetime | None = None
    completed_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    attempt_number: int = 0

    def to_dict(self) -> dict[str, Any]:
        return {
            "task_id": self.task_id,
            "status": self.status,
            "summary": self.summary,
            "artifacts": [a.to_dict() if isinstance(a, Artifact) else a for a in self.artifacts],
            "metrics": self.metrics,
            "gpu_info": self.gpu_info,
            "cost": self.cost,
            "session_id": self.session_id,
            "started_at": self.started_at.isoformat() if self.started_at else None,
            "completed_at": self.completed_at.isoformat(),
            "attempt_number": self.attempt_number,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Result:
        raw_artifacts = d.get("artifacts", [])
        artifacts = [
            Artifact.from_dict(a) if isinstance(a, dict) else a
            for a in raw_artifacts
        ]
        return cls(
            task_id=d["task_id"],
            status=d["status"],
            summary=d["summary"],
            artifacts=artifacts,
            metrics=d.get("metrics", {}),
            gpu_info=d.get("gpu_info"),
            cost=d.get("cost", _default_cost()),
            session_id=d.get("session_id"),
            started_at=datetime.fromisoformat(d["started_at"]) if d.get("started_at") else None,
            completed_at=datetime.fromisoformat(d["completed_at"]),
            attempt_number=d.get("attempt_number", 0),
        )


# ---------------------------------------------------------------------------
# Spoke (remote agent node)
# ---------------------------------------------------------------------------

@dataclass
class Spoke:
    """A remote agent node that can execute tasks."""
    id: str  # unique spoke identifier, e.g. "mac-local", "gpu-a100"
    capabilities: list[str] = field(default_factory=lambda: ["claude", "codex"])
    hostname: str = ""
    status: str = "online"  # online, offline, busy
    last_heartbeat: datetime = field(default_factory=lambda: datetime.now(UTC))
    registered_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    running_tasks: int = 0
    system_info: dict[str, Any] = field(default_factory=dict)  # cpu, gpu, disk, etc.

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "capabilities": self.capabilities,
            "hostname": self.hostname,
            "status": self.status,
            "last_heartbeat": self.last_heartbeat.isoformat(),
            "registered_at": self.registered_at.isoformat(),
            "running_tasks": self.running_tasks,
            "system_info": self.system_info,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Spoke:
        return cls(
            id=d["id"],
            capabilities=d.get("capabilities", ["claude", "codex"]),
            hostname=d.get("hostname", ""),
            status=d.get("status", "online"),
            last_heartbeat=(
                datetime.fromisoformat(d["last_heartbeat"])
                if d.get("last_heartbeat") else datetime.now(UTC)
            ),
            registered_at=(
                datetime.fromisoformat(d["registered_at"])
                if d.get("registered_at") else datetime.now(UTC)
            ),
            running_tasks=d.get("running_tasks", 0),
            system_info=d.get("system_info", {}),
        )


# ---------------------------------------------------------------------------
# Todo
# ---------------------------------------------------------------------------

def _generate_todo_id() -> str:
    date_str = datetime.now(UTC).strftime("%Y%m%d")
    suffix = secrets.token_hex(4)  # 8-char hex
    return f"todo-{date_str}-{suffix}"


@dataclass
class Todo:
    """A lightweight tracking item (not a dispatched task)."""
    description: str
    id: str = field(default_factory=_generate_todo_id)
    project: str | None = None
    due_date: datetime | None = None
    priority: Priority = Priority.NORMAL
    status: TodoStatus = TodoStatus.OPEN
    tags: list[str] = field(default_factory=list)
    source_task_id: str | None = None
    reminder_cron: str | None = None
    created_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    updated_at: datetime = field(default_factory=lambda: datetime.now(UTC))
    completed_at: datetime | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "id": self.id,
            "project": self.project,
            "description": self.description,
            "due_date": self.due_date.isoformat() if self.due_date else None,
            "priority": self.priority.value,
            "status": self.status.value,
            "tags": self.tags,
            "source_task_id": self.source_task_id,
            "reminder_cron": self.reminder_cron,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "completed_at": self.completed_at.isoformat() if self.completed_at else None,
        }

    @classmethod
    def from_dict(cls, d: dict[str, Any]) -> Todo:
        return cls(
            id=d["id"],
            project=d.get("project"),
            description=d["description"],
            due_date=datetime.fromisoformat(d["due_date"]) if d.get("due_date") else None,
            priority=Priority(d.get("priority", "normal")),
            status=TodoStatus(d.get("status", "open")),
            tags=d.get("tags", []),
            source_task_id=d.get("source_task_id"),
            reminder_cron=d.get("reminder_cron"),
            created_at=datetime.fromisoformat(d["created_at"]),
            updated_at=datetime.fromisoformat(d["updated_at"]),
            completed_at=(
                datetime.fromisoformat(d["completed_at"]) if d.get("completed_at") else None
            ),
        )


# ---------------------------------------------------------------------------
# Daily Planner
# ---------------------------------------------------------------------------

@dataclass
class TimeBlock:
    """A scheduled time block in the daily plan."""
    start: datetime
    end: datetime
    block_type: str  # "focus", "meeting", "break", "lunch", "buffer"
    title: str
    project: str | None = None
    source: str = "planner"  # "calendar", "planner", "manual"
    color_id: str | None = None  # Google Calendar color ID (1-11)
    calendar_event_id: str | None = None  # Created GCal event ID

    @property
    def duration_minutes(self) -> int:
        return int((self.end - self.start).total_seconds() / 60)

    def to_dict(self) -> dict:
        return {
            "start": self.start.isoformat(),
            "end": self.end.isoformat(),
            "block_type": self.block_type,
            "title": self.title,
            "project": self.project,
            "source": self.source,
            "duration_minutes": self.duration_minutes,
            "color_id": self.color_id,
            "calendar_event_id": self.calendar_event_id,
        }

    @classmethod
    def from_dict(cls, d: dict) -> TimeBlock:
        return cls(
            start=datetime.fromisoformat(d["start"]),
            end=datetime.fromisoformat(d["end"]),
            block_type=d["block_type"],
            title=d["title"],
            project=d.get("project"),
            source=d.get("source", "planner"),
            color_id=d.get("color_id"),
            calendar_event_id=d.get("calendar_event_id"),
        )


@dataclass
class DailyPlan:
    """A generated daily plan with time blocks."""
    date: str  # YYYY-MM-DD
    timezone: str
    blocks: list[TimeBlock]
    calendar_events: list[dict]  # Raw calendar events that were considered
    id: str = field(default_factory=lambda: f"plan-{secrets.token_hex(4)}")
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone_utc))
    briefing_message_ts: str | None = None  # Slack message ts

    def to_dict(self) -> dict:
        return {
            "id": self.id,
            "date": self.date,
            "timezone": self.timezone,
            "blocks": [b.to_dict() for b in self.blocks],
            "calendar_events": self.calendar_events,
            "created_at": self.created_at.isoformat(),
            "briefing_message_ts": self.briefing_message_ts,
        }

    @classmethod
    def from_dict(cls, d: dict) -> DailyPlan:
        return cls(
            id=d["id"],
            date=d["date"],
            timezone=d["timezone"],
            blocks=[TimeBlock.from_dict(b) for b in d.get("blocks", [])],
            calendar_events=d.get("calendar_events", []),
            created_at=(
                datetime.fromisoformat(d["created_at"])
                if d.get("created_at") else datetime.now(timezone_utc)
            ),
            briefing_message_ts=d.get("briefing_message_ts"),
        )
