"""Assessment state management.

Manages the upgrade-assessment-state.json file that tracks progress,
enables resume, and stores auto-scaling hints and calibration data.
"""

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

STATE_FILENAME = "upgrade-assessment-state.json"


def _default_state(customer: str, target_version: str, org_alias: str) -> dict:
    return {
        "customer": customer,
        "target_version": target_version,
        "org_alias": org_alias,
        "started_at": datetime.now(timezone.utc).isoformat(),
        "completed_steps": [],
        "current_step": None,
        "failed_steps": {},
        "auto_scale_hints": {},
        "calibration": {
            "packages": 0,
            "breaking_with_impact": 0,
            "estimated_days": None,
            "actual_days": None,
        },
    }


def create_state(
    working_dir: Path,
    customer: str,
    target_version: str,
    org_alias: str,
) -> dict:
    """Create a new state file. Returns the initial state."""
    state = _default_state(customer, target_version, org_alias)
    path = working_dir / STATE_FILENAME
    path.write_text(json.dumps(state, indent=2))
    logger.info("Created state file: %s", path)
    return state


def load_state(working_dir: Path) -> dict | None:
    """Load existing state file. Returns None if not found."""
    path = working_dir / STATE_FILENAME
    if not path.exists():
        return None
    return json.loads(path.read_text())


def update_state(working_dir: Path, **updates: Any) -> dict:
    """Update specific fields in the state file."""
    path = working_dir / STATE_FILENAME
    state = json.loads(path.read_text())
    state.update(updates)
    path.write_text(json.dumps(state, indent=2))
    return state


def mark_step_started(working_dir: Path, step: str) -> dict:
    """Mark a step as currently in progress."""
    return update_state(working_dir, current_step=step)


def mark_step_completed(working_dir: Path, step: str) -> dict:
    """Mark a step as completed and clear current_step."""
    state = load_state(working_dir)
    if state is None:
        raise FileNotFoundError(f"No state file in {working_dir}")
    completed = state.get("completed_steps", [])
    if step not in completed:
        completed.append(step)
    state["completed_steps"] = completed
    state["current_step"] = None
    path = working_dir / STATE_FILENAME
    path.write_text(json.dumps(state, indent=2))
    return state


def mark_step_failed(working_dir: Path, step: str, error: str) -> dict:
    """Record a step failure."""
    state = load_state(working_dir)
    if state is None:
        raise FileNotFoundError(f"No state file in {working_dir}")
    state["failed_steps"][step] = {
        "error": error,
        "timestamp": datetime.now(timezone.utc).isoformat(),
    }
    state["current_step"] = None
    path = working_dir / STATE_FILENAME
    path.write_text(json.dumps(state, indent=2))
    return state


def is_step_completed(working_dir: Path, step: str) -> bool:
    """Check if a step has already been completed."""
    state = load_state(working_dir)
    if state is None:
        return False
    return step in state.get("completed_steps", [])
