from __future__ import annotations

import json
import os
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 ls_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 _fmt(value: Any) -> str:
    return "-" if value is None else str(value)


def _format_row(row: dict[str, Any], *, long: bool) -> dict[str, str]:
    if long:
        return {
            "run_id": row["run_id"],
            "experiment": row["experiment"],
            "state": row["state"],
            "backend": row["backend"],
            "timestamp": row["timestamp"],
            "tags": ",".join(row["tags"]) or "-",
            "started_at": row["started_at"],
            "ended_at": row["ended_at"],
            "exit_code": row["exit_code"],
            "message": row["message"],
        }
    return {
        "run_id": row["run_id"],
        "experiment": row["experiment"],
        "state": row["state"],
        "backend": row["backend"],
        "timestamp": row["timestamp"],
        "tags": ",".join(row["tags"]) or "-",
    }


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 _parse_timestamp_sort_key(value: Any) -> float:
    if not value:
        return float("-inf")
    text = str(value).strip()
    if not text:
        return float("-inf")
    if text.endswith("Z"):
        text = text[:-1] + "+00:00"
    try:
        return datetime.fromisoformat(text).timestamp()
    except ValueError:
        return float("-inf")


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 list_runs(
    *,
    state_filter: str | None = None,
    long: bool = False,
    tags: list[str] | None = None,
    group_by_tag: bool = False,
    json_output: bool = False,
) -> int:
    project_root = Path.cwd().resolve()
    runs_root = project_root / DEFAULT_RUNS_DIR
    if not runs_root.exists():
        if json_output:
            print(json.dumps(ls_result_payload(items=[])))
        else:
            print("no runs found")
        return 0

    run_dirs = sorted([p for p in runs_root.iterdir() if p.is_dir()], reverse=True)
    if not run_dirs:
        if json_output:
            print(json.dumps(ls_result_payload(items=[])))
        else:
            print("no runs found")
        return 0

    selected_tags = sorted({tag.strip() for tag in (tags or []) if tag and tag.strip()})
    use_color = should_use_color()
    rows_with_sort: list[tuple[float, dict[str, Any]]] = []
    matched = 0
    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 state_filter and status.state != state_filter:
            continue
        run_tags = _normalize_tags(metadata.get("tags"))
        if not _matches_tags(run_tags, selected_tags):
            continue

        matched += 1
        rows_with_sort.append(
            (
                _parse_timestamp_sort_key(metadata.get("timestamp")),
                {
                    "run_id": run_dir.name,
                    "state": status.state,
                    "backend": metadata.get("backend_type", "?"),
                    "timestamp": iso_short(metadata.get("timestamp")),
                    "started_at": iso_short(status.started_at),
                    "ended_at": iso_short(status.ended_at),
                    "exit_code": _fmt(status.exit_code),
                    "message": _fmt(status.message),
                    "experiment": metadata.get("experiment_name", "?"),
                    "experiment_color": metadata.get("experiment_color"),
                    "tags": run_tags,
                },
            )
        )
    rows_with_sort.sort(key=lambda item: item[0], reverse=True)
    rows = [row for _, row in rows_with_sort]

    if matched == 0:
        if json_output:
            print(json.dumps(ls_result_payload(items=[])))
        else:
            print("no runs matched filter")
        return 0

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

    columns = (
        ["run_id", "experiment", "state", "backend", "timestamp", "tags", "started_at", "ended_at", "exit_code", "message"]
        if long
        else ["run_id", "experiment", "state", "backend", "timestamp", "tags"]
    )

    if group_by_tag:
        groups: dict[str, list[dict[str, Any]]] = {}
        for row in rows:
            row_tags = row["tags"] if row["tags"] else ["untagged"]
            for tag_name in row_tags:
                groups.setdefault(tag_name, []).append(row)

        for tag_name in sorted(groups):
            print(f"tag={tag_name}")
            display_rows = [_format_row(row, long=long) for row in groups[tag_name]]
            lines = render_table(columns, display_rows)
            print(lines[0])
            for idx, line in enumerate(lines):
                if idx == 0:
                    continue
                row = groups[tag_name][idx - 1]
                print(colorize(line, row["experiment_color"], enabled=use_color))
    else:
        display_rows = [_format_row(row, long=long) for row in rows]
        lines = render_table(columns, display_rows)
        print(lines[0])
        for idx, line in enumerate(lines[1:]):
            print(colorize(line, rows[idx]["experiment_color"], enabled=use_color))
    return 0
