"""Assessment state management.

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

Tools pass working_directory from arguments (the AI silently provides
the workspace path). Falls back to Path.cwd() if not provided.
"""

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"

_working_dir_override: Path | None = None


def set_working_dir(path: str | Path | None) -> None:
    """Set the working directory override from a tool argument.

    Called by tools that receive working_directory from the AI.
    Once set, all subsequent get_working_dir() calls return this path.
    """
    global _working_dir_override
    if path:
        _working_dir_override = Path(path)


def get_working_dir() -> Path:
    """Return the working directory for state and output files.

    Priority:
    1. Explicit override set by set_working_dir() (from AI argument)
    2. Path.cwd() fallback
    """
    if _working_dir_override:
        return _working_dir_override
    return Path.cwd()


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 ensure_state(**fields: Any) -> dict:
    """Create or update the state file in the working directory.

    If the state file exists, merges the given fields into it.
    If it does not exist, creates a new state file with only
    the given fields (plus an empty completed_steps list).

    This is the primary entry point for tools that need to
    persist data without knowing whether the state file exists.
    """
    working_dir = get_working_dir()
    path = working_dir / STATE_FILENAME
    if path.exists():
        current = json.loads(path.read_text())
        current.update(fields)
        path.write_text(json.dumps(current, indent=2))
        return current

    new_state: dict[str, Any] = {
        "completed_steps": [],
        "current_step": None,
        "failed_steps": {},
    }
    new_state.update(fields)
    path.write_text(json.dumps(new_state, indent=2))
    logger.info("Created state file: %s", path)
    return new_state


def create_state(
    working_dir: Path,
    customer: str,
    target_version: str,
    org_alias: str,
) -> dict:
    """Create a new state file. Returns the initial state.

    Preserves any existing fields (e.g. LMA data) if the file
    already exists, then overlays the assessment fields.
    """
    path = working_dir / STATE_FILENAME
    if path.exists():
        existing = json.loads(path.read_text())
        existing.update({
            "customer": customer,
            "target_version": target_version,
            "org_alias": org_alias,
            "started_at": datetime.now(timezone.utc).isoformat(),
            "completed_steps": existing.get("completed_steps", []),
            "current_step": None,
            "failed_steps": existing.get("failed_steps", {}),
            "auto_scale_hints": existing.get("auto_scale_hints", {}),
            "calibration": existing.get("calibration", {
                "packages": 0,
                "breaking_with_impact": 0,
                "estimated_days": None,
                "actual_days": None,
            }),
        })
        path.write_text(json.dumps(existing, indent=2))
        logger.info("Updated state file with assessment fields: %s", path)
        return existing

    state = _default_state(customer, target_version, org_alias)
    path.write_text(json.dumps(state, indent=2))
    logger.info("Created state file: %s", path)
    return state


def load_state(working_dir: Path | None = None) -> dict | None:
    """Load existing state file. Returns None if not found.

    If working_dir is None, uses get_working_dir().
    """
    if working_dir is None:
        working_dir = get_working_dir()
    path = working_dir / STATE_FILENAME
    if not path.exists():
        return None
    return json.loads(path.read_text())


def update_state(working_dir: Path | None = None, **updates: Any) -> dict:
    """Update specific fields in the state file.

    If working_dir is None, uses get_working_dir().
    """
    if working_dir is None:
        working_dir = get_working_dir()
    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 | None = None, 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 | None = None, step: str = "") -> dict:
    """Mark a step as completed and clear current_step."""
    if working_dir is None:
        working_dir = get_working_dir()
    s = load_state(working_dir)
    if s is None:
        raise FileNotFoundError(f"No state file in {working_dir}")
    completed = s.get("completed_steps", [])
    if step not in completed:
        completed.append(step)
    s["completed_steps"] = completed
    s["current_step"] = None
    path = working_dir / STATE_FILENAME
    path.write_text(json.dumps(s, indent=2))
    return s


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


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