from __future__ import annotations

import json
import os
import sys
import time
from datetime import datetime
from pathlib import Path
from typing import Any

from one_run.constants import DEFAULT_RUNS_DIR
from one_run.models import RunStatus
from one_run.payloads import ps_result_payload
from one_run.run_dir import resolve_run_paths
from one_run.utils.shell import run_capture
from one_run.status_store import read_metadata, read_status, write_status
from one_run.utils.terminal import colorize, render_table, should_use_color
from one_run.utils.timefmt import iso_short
from one_run.tracking import append_event, write_summary


def _normalize_tags(raw_tags: Any) -> list[str]:
    if not isinstance(raw_tags, list):
        return []
    return sorted({str(tag).strip() for tag in raw_tags if str(tag).strip()})


def _matches_tags(run_tags: list[str], required_tags: list[str] | None) -> bool:
    if not required_tags:
        return True
    return set(required_tags).issubset(set(run_tags))


def _is_pid_alive(pid: int) -> bool:
    try:
        os.kill(pid, 0)
        return True
    except ProcessLookupError:
        return False
    except PermissionError:
        return True


def _is_pid_zombie(pid: int) -> bool:
    try:
        proc = run_capture(["ps", "-o", "stat=", "-p", str(pid)])
    except OSError:
        return False
    if proc.returncode != 0:
        return False
    stat = (proc.stdout or "").strip()
    return "Z" in stat


def _refresh_local_status_if_possible(run_dir: Path, status: RunStatus) -> RunStatus:
    if status.state not in {"pending", "running"}:
        return status

    local_exit_code = run_dir / "local.exit_code"
    local_pid = run_dir / "local.pid"
    status_path = run_dir / "status.json"

    if local_exit_code.exists():
        try:
            exit_code = int(local_exit_code.read_text(encoding="utf-8").strip())
        except ValueError:
            exit_code = 1
        ended_at = status.ended_at or datetime.fromtimestamp(local_exit_code.stat().st_mtime)
        updated = status.model_copy(
            update={
                "state": "success" if exit_code == 0 else "failed",
                "exit_code": exit_code,
                "ended_at": ended_at,
                "message": None if exit_code == 0 else f"command exited with code {exit_code}",
            }
        )
        if updated != status:
            write_status(status_path, updated)
            resolved = resolve_run_paths(run_dir)
            append_event(resolved, run_id=run_dir.name, event="state_changed", state=updated.state, message=updated.message)
            write_summary(resolved, run_dir.name)
        return updated

    if local_pid.exists():
        try:
            pid = int(local_pid.read_text(encoding="utf-8").strip())
        except ValueError:
            pid = -1
        if pid > 0 and (_is_pid_zombie(pid) or not _is_pid_alive(pid)):
            updated = status.model_copy(
                update={
                    "state": "failed",
                    "ended_at": status.ended_at or datetime.now(),
                    "message": "process ended without exit code file",
                }
            )
            if updated != status:
                write_status(status_path, updated)
                resolved = resolve_run_paths(run_dir)
                append_event(resolved, run_id=run_dir.name, event="state_changed", state=updated.state, message=updated.message)
                write_summary(resolved, run_dir.name)
            return updated

    return status


def _collect_rows(project_root: Path, state_filter: str | None, tags: list[str] | None) -> list[dict[str, str]]:
    runs_root = project_root / DEFAULT_RUNS_DIR
    if not runs_root.exists():
        return []

    run_dirs = sorted([p for p in runs_root.iterdir() if p.is_dir()], reverse=True)
    if not run_dirs:
        return []

    selected_tags = sorted({tag.strip() for tag in (tags or []) if tag and tag.strip()})
    active_states = {"pending", "running"}
    rows_with_sort: list[tuple[float, dict[str, str]]] = []

    for run_dir in run_dirs:
        metadata_path = run_dir / "metadata.json"
        status_path = run_dir / "status.json"
        if not metadata_path.exists() or not status_path.exists():
            continue

        metadata = read_metadata(metadata_path)
        status = read_status(status_path)
        if metadata.get("backend_type") == "local":
            status = _refresh_local_status_if_possible(run_dir, status)
        if status.state not in active_states:
            continue
        if state_filter and status.state != state_filter:
            continue

        run_tags = _normalize_tags(metadata.get("tags"))
        if not _matches_tags(run_tags, selected_tags):
            continue

        backend = metadata.get("backend_type", "?")
        if backend == "local":
            pid_file = run_dir / "local.pid"
            if pid_file.exists():
                try:
                    pid = int(pid_file.read_text(encoding="utf-8").strip())
                    pid_or_job = str(pid)
                    alive = "yes" if _is_pid_alive(pid) else "no"
                except ValueError:
                    pid_or_job = "-"
                    alive = "unknown"
            else:
                pid_or_job = "-"
                alive = "unknown"
        else:
            pid_or_job = str(metadata.get("slurm_job_id") or "-")
            alive = "unknown"

        started_sort = status.started_at.timestamp() if status.started_at else float("-inf")
        rows_with_sort.append(
            (
                started_sort,
                {
                    "run_id": run_dir.name,
                    "state": status.state,
                    "backend": backend,
                    "pid": pid_or_job,
                    "alive": alive,
                    "started_at": iso_short(status.started_at),
                    "experiment": str(metadata.get("experiment_name", "?")),
                    "tags": ",".join(run_tags) or "-",
                    "experiment_color": str(metadata.get("experiment_color") or ""),
                },
            )
        )
    rows_with_sort.sort(key=lambda item: item[0], reverse=True)
    return [row for _, row in rows_with_sort]


def _build_lines(rows: list[dict[str, str]], *, use_color: bool) -> list[str]:
    if not rows:
        return ["no active runs matched filter"]
    columns = ["run_id", "state", "backend", "pid", "alive", "started_at", "experiment", "tags"]
    lines = render_table(columns, rows)
    if not use_color:
        return lines
    colored = [lines[0]]
    for idx, line in enumerate(lines[1:]):
        colored.append(colorize(line, rows[idx].get("experiment_color"), enabled=True))
    return colored


def _print_lines(lines: list[str]) -> None:
    for line in lines:
        print(line)


def _redraw_lines_tty(lines: list[str], previous_count: int) -> int:
    if previous_count > 0:
        for _ in range(previous_count):
            sys.stdout.write("\x1b[1A")
            sys.stdout.write("\x1b[2K")
    for line in lines:
        sys.stdout.write(line + "\n")
    sys.stdout.flush()
    return len(lines)


def list_processes(
    *,
    state_filter: str | None = None,
    tags: list[str] | None = None,
    json_output: bool = False,
    watch: bool = False,
    interval_sec: float = 2.0,
) -> int:
    project_root = Path.cwd().resolve()
    use_color = should_use_color()

    rows = _collect_rows(project_root, state_filter, tags)
    if not rows and json_output:
        print(json.dumps(ps_result_payload(items=[])))
        return 0

    if watch and json_output:
        watch = False

    if json_output:
        print(json.dumps(ps_result_payload(items=rows)))
        return 0

    if not watch:
        _print_lines(_build_lines(rows, use_color=use_color))
        return 0

    if not sys.stdout.isatty():
        _print_lines(_build_lines(rows, use_color=use_color))
        return 0

    prev_count = 0
    sys.stdout.write("\x1b[?25l")
    sys.stdout.flush()
    try:
        while True:
            current_rows = _collect_rows(project_root, state_filter, tags)
            lines = _build_lines(current_rows, use_color=use_color)
            prev_count = _redraw_lines_tty(lines, prev_count)
            time.sleep(max(interval_sec, 0.1))
    except KeyboardInterrupt:
        pass
    finally:
        sys.stdout.write("\x1b[?25h")
        sys.stdout.flush()
        return 0
