"""Core data models for ABREKA."""

from __future__ import annotations

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


class ExperimentStatus(Enum):
    PROPOSING = "proposing"
    PROPOSED = "proposed"
    IMPLEMENTING = "implementing"
    TESTING = "testing"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"


class BacklogStatus(Enum):
    OPEN = "open"
    RESOLVED = "resolved"
    DISMISSED = "dismissed"


# Valid status transitions
VALID_TRANSITIONS: dict[ExperimentStatus, set[ExperimentStatus]] = {
    ExperimentStatus.PROPOSING: {ExperimentStatus.PROPOSED, ExperimentStatus.FAILED},
    ExperimentStatus.PROPOSED: {ExperimentStatus.IMPLEMENTING, ExperimentStatus.FAILED},
    ExperimentStatus.IMPLEMENTING: {ExperimentStatus.TESTING, ExperimentStatus.FAILED},
    ExperimentStatus.TESTING: {ExperimentStatus.RUNNING, ExperimentStatus.FAILED},
    ExperimentStatus.RUNNING: {ExperimentStatus.COMPLETED, ExperimentStatus.FAILED},
    ExperimentStatus.COMPLETED: set(),
    ExperimentStatus.FAILED: set(),
}


@dataclass
class Experiment:
    id: str
    status: ExperimentStatus = ExperimentStatus.PROPOSED
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    completed_at: datetime | None = None
    hypothesis: str = ""
    method: str = ""
    parents: list[str] = field(default_factory=list)
    tags: list[str] = field(default_factory=list)
    metrics: dict[str, dict[str, float]] = field(default_factory=dict)
    runtime_costs: dict[str, dict[str, object]] = field(default_factory=dict)
    failures: list[dict[str, object]] = field(default_factory=list)
    failure_reason: str = ""

    def to_dict(self) -> dict:
        d: dict = {
            "id": self.id,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "hypothesis": self.hypothesis,
            "method": self.method,
            "parents": self.parents,
            "tags": self.tags,
        }
        if self.completed_at:
            d["completed_at"] = self.completed_at.isoformat()
        if self.metrics:
            d["metrics"] = self.metrics
        if self.runtime_costs:
            d["runtime_costs"] = self.runtime_costs
        if self.failures:
            d["failures"] = self.failures
        if self.failure_reason:
            d["failure"] = {"reason": self.failure_reason}
        return d

    @classmethod
    def from_dict(cls, d: dict) -> Experiment:
        failure = d.get("failure", {})
        completed_at = d.get("completed_at")
        return cls(
            id=d["id"],
            status=ExperimentStatus(d["status"]),
            created_at=datetime.fromisoformat(d["created_at"]),
            completed_at=datetime.fromisoformat(completed_at) if completed_at else None,
            hypothesis=d.get("hypothesis", ""),
            method=d.get("method", ""),
            parents=d.get("parents", []),
            tags=d.get("tags", []),
            metrics=d.get("metrics", {}),
            runtime_costs=d.get("runtime_costs", {}),
            failures=d.get("failures", []),
            failure_reason=failure.get("reason", "") if isinstance(failure, dict) else "",
        )

    def summary(self) -> dict:
        """Lightweight dict for the index."""
        s: dict = {
            "id": self.id,
            "status": self.status.value,
            "hypothesis": self.hypothesis,
            "method": self.method,
            "parents": self.parents,
            "tags": self.tags,
            "created_at": self.created_at.isoformat(),
        }
        if self.metrics:
            s["metrics"] = self.metrics
        if self.runtime_costs:
            total_runtime = 0.0
            for stage in self.runtime_costs.values():
                if not isinstance(stage, dict):
                    continue
                try:
                    total_runtime += float(stage.get("seconds_total", 0.0))
                except (TypeError, ValueError):
                    continue
            s["runtime_seconds_total"] = total_runtime
        if self.failures:
            s["failure_count"] = len(self.failures)
        return s


@dataclass
class ProjectConfig:
    name: str = ""
    lib: str = ""
    test_command: str = "uv run pytest"


@dataclass
class ExperimentConfig:
    primary_metric: str = "val:accuracy"
    metric_direction: str = "maximize"
    run_template: str = "./run.py"
    explore_probability: float = 0.0


@dataclass
class RemoteRunConfig:
    host: str = ""
    project_dir: str = ""
    bootstrap_command: str = ""
    poll_seconds: int = 5
    startup_command: str = ""
    teardown_command: str = ""


@dataclass
class RunConfig:
    timeout_minutes: int = 30
    require_tests: bool = True
    mode: str = "local"
    max_repair_attempts: int = 3
    remote: RemoteRunConfig = field(default_factory=RemoteRunConfig)


@dataclass
class AgentConfig:
    provider: str = "openai-codex"
    model: str = "gpt-5.4"
    thinking_level: str = "xhigh"
    propose_model: str = ""
    implement_model: str = ""
    evaluate_model: str = ""
    session_capture_mode: str = "slim"
    pi_extra_args: list[str] = field(default_factory=list)


@dataclass
class Config:
    project: ProjectConfig = field(default_factory=ProjectConfig)
    experiment: ExperimentConfig = field(default_factory=ExperimentConfig)
    run: RunConfig = field(default_factory=RunConfig)
    agent: AgentConfig = field(default_factory=AgentConfig)

    def to_dict(self) -> dict:
        d: dict = {
            "project": {
                "name": self.project.name,
                "lib": self.project.lib,
                "test_command": self.project.test_command,
            },
            "experiment": {
                "primary_metric": self.experiment.primary_metric,
                "metric_direction": self.experiment.metric_direction,
                "run_template": self.experiment.run_template,
                **({"explore_probability": self.experiment.explore_probability}
                   if self.experiment.explore_probability > 0 else {}),
            },
            "run": {
                "timeout_minutes": self.run.timeout_minutes,
                "require_tests": self.run.require_tests,
                "mode": self.run.mode,
                "max_repair_attempts": self.run.max_repair_attempts,
            },
            "agent": {
                "provider": self.agent.provider,
                "model": self.agent.model,
                "thinking_level": self.agent.thinking_level,
            },
        }
        if (
            self.run.mode == "remote"
            or self.run.remote.host
            or self.run.remote.project_dir
            or self.run.remote.bootstrap_command
            or self.run.remote.poll_seconds != 5
            or self.run.remote.startup_command
            or self.run.remote.teardown_command
        ):
            d["run"]["remote"] = {
                "host": self.run.remote.host,
                "project_dir": self.run.remote.project_dir,
                **({"bootstrap_command": self.run.remote.bootstrap_command}
                   if self.run.remote.bootstrap_command else {}),
                **({"poll_seconds": self.run.remote.poll_seconds}
                   if self.run.remote.poll_seconds != 5 else {}),
                **({"startup_command": self.run.remote.startup_command}
                   if self.run.remote.startup_command else {}),
                **({"teardown_command": self.run.remote.teardown_command}
                   if self.run.remote.teardown_command else {}),
            }
        if self.agent.propose_model:
            d["agent"]["propose_model"] = self.agent.propose_model
        if self.agent.implement_model:
            d["agent"]["implement_model"] = self.agent.implement_model
        if self.agent.evaluate_model:
            d["agent"]["evaluate_model"] = self.agent.evaluate_model
        if self.agent.session_capture_mode != "slim":
            d["agent"]["session_capture_mode"] = self.agent.session_capture_mode
        if self.agent.pi_extra_args:
            d["agent"]["pi"] = {"extra_args": self.agent.pi_extra_args}
        return d

    @classmethod
    def from_dict(cls, d: dict) -> Config:
        proj = d.get("project", {})
        exp = d.get("experiment", {})
        run = d.get("run", {})
        remote = run.get("remote", {})
        ag = d.get("agent", {})
        pi = ag.get("pi", {})
        return cls(
            project=ProjectConfig(
                name=proj.get("name", ""),
                lib=proj.get("lib", ""),
                test_command=proj.get("test_command", "uv run pytest"),
            ),
            experiment=ExperimentConfig(
                primary_metric=exp.get("primary_metric", "val:accuracy"),
                metric_direction=exp.get("metric_direction", "maximize"),
                run_template=exp.get("run_template", "./run.py"),
                explore_probability=exp.get("explore_probability", 0.0),
            ),
            run=RunConfig(
                timeout_minutes=run.get("timeout_minutes", 30),
                require_tests=run.get("require_tests", True),
                mode=run.get("mode", "local"),
                max_repair_attempts=run.get("max_repair_attempts", 3),
                remote=RemoteRunConfig(
                    host=remote.get("host", ""),
                    project_dir=remote.get("project_dir", ""),
                    bootstrap_command=remote.get("bootstrap_command", ""),
                    poll_seconds=remote.get("poll_seconds", 5),
                    startup_command=remote.get("startup_command", ""),
                    teardown_command=remote.get("teardown_command", ""),
                ),
            ),
            agent=AgentConfig(
                provider=ag.get("provider", "openai-codex"),
                model=ag.get("model", "gpt-5.4"),
                thinking_level=ag.get("thinking_level", "xhigh"),
                propose_model=ag.get("propose_model", ""),
                implement_model=ag.get("implement_model", ""),
                evaluate_model=ag.get("evaluate_model", ""),
                session_capture_mode=ag.get("session_capture_mode", "slim"),
                pi_extra_args=pi.get("extra_args", []),
            ),
        )


@dataclass
class WikiPage:
    """A project wiki page: durable cross-experiment knowledge.

    Slug-keyed, markdown-bodied. Source of truth is the file at
    ``abreka/wiki/pages/<slug>.md`` (TOML frontmatter + markdown body).
    """

    slug: str
    title: str = ""
    summary: str = ""
    tags: list[str] = field(default_factory=list)
    body: str = ""
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    created_by_exp: str = ""
    updated_by_exp: str = ""

    def index_entry(self) -> dict:
        """Lightweight dict for the wiki index."""
        return {
            "slug": self.slug,
            "title": self.title,
            "summary": self.summary,
            "tags": list(self.tags),
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
            "created_by_exp": self.created_by_exp,
            "updated_by_exp": self.updated_by_exp,
        }


@dataclass
class BacklogItem:
    id: str
    title: str = ""
    description: str = ""
    tags: list[str] = field(default_factory=list)
    priority: int = 0
    source_exp: str = ""
    resolved_exp: str = ""
    status: BacklogStatus = BacklogStatus.OPEN
    dismiss_reason: str = ""
    created_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
    updated_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))

    def to_dict(self) -> dict:
        d: dict = {
            "id": self.id,
            "title": self.title,
            "description": self.description,
            "tags": self.tags,
            "priority": self.priority,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }
        if self.source_exp:
            d["source_exp"] = self.source_exp
        if self.resolved_exp:
            d["resolved_exp"] = self.resolved_exp
        if self.dismiss_reason:
            d["dismiss_reason"] = self.dismiss_reason
        return d

    @classmethod
    def from_dict(cls, d: dict) -> BacklogItem:
        return cls(
            id=d["id"],
            title=d.get("title", ""),
            description=d.get("description", ""),
            tags=d.get("tags", []),
            priority=d.get("priority", 0),
            source_exp=d.get("source_exp", ""),
            resolved_exp=d.get("resolved_exp", ""),
            status=BacklogStatus(d.get("status", "open")),
            dismiss_reason=d.get("dismiss_reason", ""),
            created_at=datetime.fromisoformat(d["created_at"]),
            updated_at=datetime.fromisoformat(d["updated_at"]),
        )

    def summary(self) -> dict:
        """Lightweight dict for the backlog index."""
        return {
            "id": self.id,
            "title": self.title,
            "description": self.description,
            "tags": self.tags,
            "priority": self.priority,
            "source_exp": self.source_exp,
            "resolved_exp": self.resolved_exp,
            "status": self.status.value,
            "created_at": self.created_at.isoformat(),
            "updated_at": self.updated_at.isoformat(),
        }
