"""Repair software-job outputs written under the duplicated legacy user path.

The command is dry-run by default. It reads its database URL only from the
explicit ``ZCBOT_MIGRATION_DB_URL`` environment variable and never loads .env.
Run schema migration 0033 before applying this repair.
"""
# ruff: noqa: I001
from __future__ import annotations

import argparse
import os
import sys
from dataclasses import dataclass
from pathlib import Path
from uuid import UUID

from sqlalchemy import create_engine, select
from sqlalchemy.orm import Session
from sqlalchemy.orm.attributes import flag_modified

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from core.software_contracts import DEFAULT_CAPABILITIES, get_contract
from core.software_jobs import software_job_output_path
from core.storage.models import Artifact, SoftwareJob, Task
from core.paths import from_db_path


@dataclass(frozen=True)
class RepairPlan:
    job_id: UUID
    user_root: Path
    working_dir: Path
    legacy_dir: Path
    output_dir: Path
    move_legacy_dir: bool
    metadata_moves: tuple[tuple[Path, Path], ...]


def _user_root(working_dir: Path, user_id: UUID) -> Path:
    expected = str(user_id)
    for candidate in (working_dir, *working_dir.parents):
        if candidate.name == expected and candidate.parent.name == "users":
            return candidate.resolve()
    raise RuntimeError(f"task working_dir has no user root for {user_id}")


def _within(root: Path, target: Path) -> Path:
    resolved = target.resolve()
    resolved.relative_to(root.resolve())
    return resolved


def build_plan(job: SoftwareJob, task: Task) -> RepairPlan | None:
    capability = getattr(job, "capability", DEFAULT_CAPABILITIES[0])
    contract = get_contract(capability)
    working_dir = from_db_path(task.working_dir).resolve()
    user_root = _user_root(working_dir, job.user_id)
    output_dir = _within(
        user_root, working_dir / contract.output_namespace / str(job.job_id)
    )
    legacy_dir = _within(
        user_root,
        user_root / Path(task.working_dir) / contract.output_namespace / str(job.job_id),
    )
    legacy_exists = legacy_dir.is_dir()
    output_exists = output_dir.is_dir()
    if legacy_exists and output_exists:
        raise RuntimeError(f"job {job.job_id}: legacy and target output directories both exist")
    if not legacy_exists and not output_exists:
        print(f"[WARN] job {job.job_id}: output directory is missing")
        return None
    source_dir = legacy_dir if legacy_exists else output_dir
    metadata_moves: list[tuple[Path, Path]] = []
    for item in job.artifact_manifest or []:
        output_id = str((item or {}).get("source_artifact_id") or "")
        if not output_id or contract.output_spec(output_id).publish:
            continue
        source = source_dir / str((item or {}).get("filename") or "")
        destination = source_dir / software_job_output_path(capability, output_id)
        if source == destination or not source.exists():
            continue
        if destination.exists():
            raise RuntimeError(f"job {job.job_id}: metadata destination already exists")
        metadata_moves.append((source, destination))
    return RepairPlan(
        job_id=job.job_id,
        user_root=user_root,
        working_dir=working_dir,
        legacy_dir=legacy_dir,
        output_dir=output_dir,
        move_legacy_dir=legacy_exists,
        metadata_moves=tuple(metadata_moves),
    )


def apply_files(plan: RepairPlan) -> None:
    for _, destination in plan.metadata_moves:
        destination.parent.mkdir(parents=True, exist_ok=True)
    for source, destination in plan.metadata_moves:
        os.replace(source, destination)
    if plan.move_legacy_dir:
        plan.output_dir.parent.mkdir(parents=True, exist_ok=True)
        os.replace(plan.legacy_dir, plan.output_dir)


def update_rows(session: Session, job: SoftwareJob, task: Task, plan: RepairPlan) -> int:
    capability = getattr(job, "capability", DEFAULT_CAPABILITIES[0])
    contract = get_contract(capability)
    artifact_count = 0
    manifest = [dict(item) for item in (job.artifact_manifest or [])]
    task_prefix = plan.working_dir.relative_to(plan.user_root).as_posix()
    for item in manifest:
        output_id = str(item.get("source_artifact_id") or "")
        if not output_id:
            continue
        path = (
            f"{contract.output_namespace}/{job.job_id}/"
            f"{software_job_output_path(capability, output_id)}"
        )
        item["path"] = path
        raw_artifact_id = item.get("artifact_id")
        if not raw_artifact_id:
            continue
        try:
            artifact_id = UUID(str(raw_artifact_id))
        except ValueError:
            continue
        artifact = session.get(Artifact, artifact_id)
        if artifact is None or artifact.user_id != job.user_id:
            raise RuntimeError(f"job {job.job_id}: artifact {artifact_id} is missing")
        artifact.current_path = f"{task_prefix}/{path}"
        if contract.output_spec(output_id).publish:
            artifact.software_job_id = job.job_id
        artifact_count += 1
    job.artifact_manifest = manifest
    flag_modified(job, "artifact_manifest")
    return artifact_count


def repair(session: Session, *, apply: bool, job_id: UUID | None = None) -> tuple[int, int]:
    statement = (
        select(SoftwareJob, Task)
        .join(Task, Task.task_id == SoftwareJob.task_id)
        .where(SoftwareJob.status == "succeeded")
        .order_by(SoftwareJob.created_at, SoftwareJob.job_id)
    )
    if job_id is not None:
        statement = statement.where(SoftwareJob.job_id == job_id)
    rows = session.execute(statement).all()
    prepared: list[tuple[SoftwareJob, Task, RepairPlan]] = []
    for job, task in rows:
        plan = build_plan(job, task)
        if plan is None:
            continue
        print(
            f"[INFO] job={job.job_id} move_dir={plan.move_legacy_dir} "
            f"metadata_moves={len(plan.metadata_moves)} target={plan.output_dir}"
        )
        prepared.append((job, task, plan))
    if not apply:
        return len(prepared), 0
    updated_artifacts = 0
    for job, task, plan in prepared:
        apply_files(plan)
        updated_artifacts += update_rows(session, job, task, plan)
    session.flush()
    return len(prepared), updated_artifacts


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--apply", action="store_true", help="apply filesystem and database changes")
    parser.add_argument("--job-id", type=UUID, help="limit repair to one software job")
    args = parser.parse_args()
    database_url = os.environ.get("ZCBOT_MIGRATION_DB_URL", "").strip()
    if not database_url:
        print("[ERR] ZCBOT_MIGRATION_DB_URL is required", file=sys.stderr)
        return 2
    engine = create_engine(database_url, pool_pre_ping=True, future=True)
    try:
        with Session(engine, future=True) as session:
            try:
                jobs, artifacts = repair(
                    session, apply=args.apply, job_id=args.job_id
                )
                if args.apply:
                    session.commit()
                else:
                    session.rollback()
            except Exception:
                session.rollback()
                raise
    finally:
        engine.dispose()
    mode = "applied" if args.apply else "validated"
    print(f"[OK] {mode} jobs: {jobs}")
    print(f"[OK] updated artifacts: {artifacts}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())