from __future__ import annotations

import uuid
from datetime import datetime, timezone
from enum import Enum
from typing import Any

from pydantic import BaseModel, Field, model_validator


def _now() -> datetime:
    return datetime.now(timezone.utc)


def _id() -> str:
    return uuid.uuid4().hex[:12]


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

class LockState(str, Enum):
    HINT = "hint"
    PROPOSED = "proposed"
    FROZEN = "frozen"


class PhaseExecutionStatus(str, Enum):
    PENDING = "pending"
    READY = "ready"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"
    INVALIDATED = "invalidated"


class PhaseName(str, Enum):
    DISCOVERY = "discovery"
    PREPROCESSING = "preprocessing"
    FEATURIZATION = "featurization_or_materialization"
    TRAINING = "training"
    EVALUATION = "evaluation"
    ANALYSIS = "analysis"


class TaskType(str, Enum):
    CLASSIFICATION = "classification"
    REGRESSION = "regression"
    TIME_SERIES = "time_series"
    NLP = "nlp"
    VISION = "vision"
    RANKING = "ranking"
    OTHER = "other"


class ExperimentStatus(str, Enum):
    KEEP = "keep"
    DISCARD = "discard"
    VERIFY = "verify"
    CRASH = "crash"


class ActivityType(str, Enum):
    THINKING = "thinking"
    SEARCHING = "searching"
    SEARCH_RESULT = "search_result"
    EDITING = "editing"
    TRAINING = "training"
    EVALUATING = "evaluating"
    DECIDING = "deciding"
    LOGGING = "logging"
    WAITING = "waiting"
    ERROR = "error"


class ConversationState(str, Enum):
    DRAFT = "draft"
    NEEDS_INPUT = "needs_input"
    READY_TO_PREPROCESS = "ready_to_preprocess"
    READY_TO_TRAIN = "ready_to_train"
    RUNNING = "running"
    PAUSED = "paused"
    REVIEWING_RESULTS = "reviewing_results"
    COMPLETED = "completed"
    FAILED = "failed"
    STOPPED = "stopped"


class DecisionType(str, Enum):
    DATASET = "dataset"
    PREPROCESSING = "preprocessing"
    EVALUATION = "evaluation"
    MODEL = "model"
    TARGET = "target"
    RUN_PLAN = "run_plan"
    POST_RUN_RECOMMENDATION = "post_run_recommendation"
    GENERAL = "general"


class DecisionStatus(str, Enum):
    NEEDS_INPUT = "needs_input"
    PROPOSED = "proposed"
    APPROVED = "approved"
    REJECTED = "rejected"
    SUPERSEDED = "superseded"
    BLOCKED = "blocked"
    COMPLETED = "completed"


class NextActionType(str, Enum):
    ANSWER_QUESTION = "answer_question"
    APPROVE_DECISION = "approve_decision"
    REVISE_DECISION = "revise_decision"
    RUN_PREPROCESSING = "run_preprocessing"
    RUN_TRAINING = "run_training"
    REVIEW_RESULTS = "review_results"
    FINALIZE = "finalize"
    STOP = "stop"


# ---------------------------------------------------------------------------
# Assets
# ---------------------------------------------------------------------------

class DatasetAsset(BaseModel):
    path: str
    source_type: str = ""  # "local", "package", "url"
    source_identifier: str = ""  # e.g. "sklearn.datasets.load_iris", a URL, etc.
    format: str = ""
    fingerprint: str = ""
    schema_info: dict[str, Any] = Field(default_factory=dict)
    target_column: str = ""
    task_type: TaskType | None = None
    num_rows: int | None = None
    num_features: int | None = None
    lock_state: LockState = LockState.HINT


class ReferenceAsset(BaseModel):
    path: str
    asset_type: str = ""  # repo, notebook, paper, model
    extracted_info: dict[str, Any] = Field(default_factory=dict)
    lock_state: LockState = LockState.HINT


class PreprocessingPlan(BaseModel):
    source: str = ""  # inferred, copied, adapted, user-defined
    transforms: list[dict[str, Any]] = Field(default_factory=list)
    reproducibility_label: str = ""  # exact, best_effort, adapted, cannot_replicate
    lock_state: LockState = LockState.HINT


# ---------------------------------------------------------------------------
# Eval and model specs
# ---------------------------------------------------------------------------

class EvalSpec(BaseModel):
    task_type: TaskType | None = None
    primary_metric: str = ""
    secondary_metrics: list[str] = Field(default_factory=list)
    split_strategy: str = ""  # holdout, stratified, grouped, time_series, kfold
    num_seeds: int = 1
    constraints: dict[str, Any] = Field(default_factory=dict)
    lock_state: LockState = LockState.HINT


class ModelSpec(BaseModel):
    framework: str = ""
    architecture: str = ""
    hyperparameters: dict[str, Any] = Field(default_factory=dict)
    allowed_families: list[str] = Field(default_factory=list)
    search_space: dict[str, Any] = Field(default_factory=dict)
    hardware_constraints: dict[str, Any] = Field(default_factory=dict)
    lock_state: LockState = LockState.HINT


class PhaseArtifact(BaseModel):
    name: str
    path: str = ""
    artifact_type: str = ""
    fingerprint: str = ""
    created_at: datetime = Field(default_factory=_now)
    metadata: dict[str, Any] = Field(default_factory=dict)


class ContractPhase(BaseModel):
    name: PhaseName
    label: str = ""
    order: int = 0
    description: str = ""
    depends_on: list[PhaseName] = Field(default_factory=list)
    lock_state: LockState = LockState.HINT
    status: PhaseExecutionStatus = PhaseExecutionStatus.PENDING
    is_iterative: bool = False
    reusable: bool = True
    invalidated: bool = False
    invalidation_reason: str = ""
    artifact_refs: list[PhaseArtifact] = Field(default_factory=list)
    script_path: str = ""
    hardware_profile: dict[str, Any] = Field(default_factory=dict)
    metadata: dict[str, Any] = Field(default_factory=dict)
    last_run_at: datetime | None = None


class DecisionRecord(BaseModel):
    id: str = Field(default_factory=_id)
    type: DecisionType = DecisionType.GENERAL
    title: str = ""
    status: DecisionStatus = DecisionStatus.PROPOSED
    summary: str = ""
    rationale: str = ""
    evidence: list[str] = Field(default_factory=list)
    before: dict[str, Any] = Field(default_factory=dict)
    after: dict[str, Any] = Field(default_factory=dict)
    blocking_reasons: list[str] = Field(default_factory=list)
    downstream_impact: list[str] = Field(default_factory=list)
    required_user_input: list[str] = Field(default_factory=list)
    phase_impacts: list[str] = Field(default_factory=list)
    spec_updates: dict[str, Any] = Field(default_factory=dict)
    user_note: str = ""
    created_at: datetime = Field(default_factory=_now)
    updated_at: datetime = Field(default_factory=_now)


class NextAction(BaseModel):
    id: str = Field(default_factory=_id)
    action_type: NextActionType
    label: str = ""
    description: str = ""
    decision_id: str | None = None
    phase_name: str | None = None
    enabled: bool = True


class Readiness(BaseModel):
    current_stage: ConversationState = ConversationState.DRAFT
    can_run_preprocessing: bool = False
    can_run_training: bool = False
    can_finalize: bool = False
    blocking_reasons: list[str] = Field(default_factory=list)
    next_actions: list[NextAction] = Field(default_factory=list)
    reusable_artifacts: list[str] = Field(default_factory=list)
    active_decision_id: str | None = None
    pending_decision_count: int = 0
    summary: str = ""
    agent_status: str = "idle"
    agent_status_detail: str = ""


def _default_phase_map() -> dict[str, ContractPhase]:
    return {
        PhaseName.DISCOVERY.value: ContractPhase(
            name=PhaseName.DISCOVERY,
            label="Discovery",
            order=0,
            description="Inspect the dataset, references, and constraints before execution.",
            depends_on=[],
            status=PhaseExecutionStatus.READY,
            reusable=False,
        ),
        PhaseName.PREPROCESSING.value: ContractPhase(
            name=PhaseName.PREPROCESSING,
            label="Preprocessing",
            order=1,
            description="Run heavy preprocessing once and save reusable outputs.",
            depends_on=[PhaseName.DISCOVERY],
        ),
        PhaseName.FEATURIZATION.value: ContractPhase(
            name=PhaseName.FEATURIZATION,
            label="Featurization / Materialization",
            order=2,
            description="Materialize reusable intermediate artifacts for downstream runs.",
            depends_on=[PhaseName.PREPROCESSING],
        ),
        PhaseName.TRAINING.value: ContractPhase(
            name=PhaseName.TRAINING,
            label="Training",
            order=3,
            description="Iterate on model training using reusable upstream artifacts.",
            depends_on=[PhaseName.FEATURIZATION],
            is_iterative=True,
            reusable=False,
        ),
        PhaseName.EVALUATION.value: ContractPhase(
            name=PhaseName.EVALUATION,
            label="Evaluation",
            order=4,
            description="Evaluate each training iteration with the frozen evaluation setup.",
            depends_on=[PhaseName.TRAINING],
            is_iterative=True,
            reusable=False,
        ),
        PhaseName.ANALYSIS.value: ContractPhase(
            name=PhaseName.ANALYSIS,
            label="Analysis",
            order=5,
            description="Summarize learnings and decide what to change next.",
            depends_on=[PhaseName.EVALUATION],
            reusable=False,
        ),
    }


def builtin_phase_order() -> list[str]:
    return [phase.value for phase in PhaseName]


def ensure_builtin_phases(contract: "FrozenContract") -> "FrozenContract":
    defaults = _default_phase_map()
    if not contract.phase_order:
        contract.phase_order = builtin_phase_order()
    for phase_name in contract.phase_order:
        if phase_name not in contract.phases and phase_name in defaults:
            contract.phases[phase_name] = defaults[phase_name].model_copy(deep=True)
    for phase_name, phase in defaults.items():
        if phase_name not in contract.phases:
            contract.phases[phase_name] = phase.model_copy(deep=True)
    return contract


# ---------------------------------------------------------------------------
# Frozen contract
# ---------------------------------------------------------------------------

class FrozenContract(BaseModel):
    id: str = Field(default_factory=_id)
    created_at: datetime = Field(default_factory=_now)
    conversation_id: str = ""

    dataset: DatasetAsset | None = None
    reference: ReferenceAsset | None = None
    preprocessing: PreprocessingPlan | None = None
    eval_spec: EvalSpec | None = None
    model_spec: ModelSpec | None = None

    target_threshold: float | None = None
    threshold_direction: str = "higher_is_better"  # or "lower_is_better"
    budget_minutes: int | None = None
    max_experiments: int | None = None

    phase_order: list[str] = Field(default_factory=builtin_phase_order)
    phases: dict[str, ContractPhase] = Field(default_factory=_default_phase_map)
    frozen: bool = False

    @model_validator(mode="after")
    def _ensure_builtin_phases_present(self) -> "FrozenContract":
        return ensure_builtin_phases(self)


# ---------------------------------------------------------------------------
# Experiment tracking
# ---------------------------------------------------------------------------

class ExperimentResult(BaseModel):
    id: str = Field(default_factory=_id)
    contract_id: str = ""
    run_number: int = 0
    created_at: datetime = Field(default_factory=_now)

    hypothesis: str = ""
    code_diff: str = ""
    metrics: dict[str, Any] = Field(default_factory=dict)
    timing_seconds: float | None = None
    memory_peak_mb: float | None = None
    phase_name: str = PhaseName.TRAINING.value
    phase_artifacts: list[PhaseArtifact] = Field(default_factory=list)

    status: ExperimentStatus = ExperimentStatus.DISCARD
    verdict_reason: str = ""
    lessons_learned: str = ""
    next_direction: str = ""
    search_queries: list[dict[str, str]] = Field(default_factory=list)


class ExperimentJournal(BaseModel):
    contract_id: str = ""
    experiments: list[ExperimentResult] = Field(default_factory=list)
    agent_commentary: str = ""
    best_experiment_id: str | None = None
    best_metric_value: float | None = None
    phase_history: list[dict[str, Any]] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Agent activity
# ---------------------------------------------------------------------------

class AgentActivity(BaseModel):
    timestamp: datetime = Field(default_factory=_now)
    activity_type: ActivityType
    summary: str = ""
    detail: str = ""
    experiment_id: str | None = None
    phase_name: str | None = None


# ---------------------------------------------------------------------------
# Chat messages
# ---------------------------------------------------------------------------

class ChatMessage(BaseModel):
    id: str = Field(default_factory=_id)
    role: str  # "user" or "assistant"
    content: str
    timestamp: datetime = Field(default_factory=_now)
    metadata: dict[str, Any] = Field(default_factory=dict)


class Conversation(BaseModel):
    id: str = Field(default_factory=_id)
    title: str = "New Project"
    subtitle: str = ""
    workspace_root: str = ""
    autonomy_level: str = "guided"
    created_at: datetime = Field(default_factory=_now)
    updated_at: datetime = Field(default_factory=_now)
    state: ConversationState = ConversationState.DRAFT
    messages: list[ChatMessage] = Field(default_factory=list)
    contract: FrozenContract = Field(default_factory=FrozenContract)
    decisions: list[DecisionRecord] = Field(default_factory=list)
    readiness: Readiness = Field(default_factory=Readiness)



# ---------------------------------------------------------------------------
# Hardware
# ---------------------------------------------------------------------------

class HardwareProfile(BaseModel):
    cpu_count: int = 0
    ram_total_gb: float = 0.0
    gpu_available: bool = False
    gpu_type: str = ""  # cuda, mps, none
    gpu_count: int = 0
    gpu_names: list[str] = Field(default_factory=list)
    gpu_vram_gb: list[float] = Field(default_factory=list)
    detected_at: datetime = Field(default_factory=_now)
