#!/usr/bin/env python3
from __future__ import annotations

import argparse
import json
import sys
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any


def utc_now() -> datetime:
    return datetime.now(timezone.utc)


def to_iso(value: datetime) -> str:
    return value.astimezone(timezone.utc).replace(microsecond=0).isoformat()


def parse_iso(value: str | None) -> datetime | None:
    raw = (value or "").strip()
    if not raw:
        return None
    try:
        return datetime.fromisoformat(raw.replace("Z", "+00:00")).astimezone(timezone.utc)
    except ValueError:
        return None


def load_state(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
    try:
        payload = json.loads(path.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
    if not isinstance(payload, dict):
        return {"version": 1, "updated_at": to_iso(utc_now()), "allocations": []}
    payload.setdefault("version", 1)
    payload.setdefault("updated_at", to_iso(utc_now()))
    allocations = payload.get("allocations")
    if not isinstance(allocations, list):
        payload["allocations"] = []
    return payload


def save_state(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload["updated_at"] = to_iso(utc_now())
    path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")


def parse_slots(raw: str) -> list[str]:
    slots = [item.strip() for item in raw.split(",") if item.strip()]
    if not slots:
        raise ValueError("at least one slot is required")
    return slots


def render_url(slot: str, template: str | None) -> str | None:
    if not template:
        return None
    return template.replace("{slot}", slot)


def prune_expired(allocations: list[dict[str, Any]], now: datetime) -> tuple[list[dict[str, Any]], int]:
    kept: list[dict[str, Any]] = []
    removed = 0
    for item in allocations:
        expires = parse_iso(str(item.get("expires_at") or ""))
        if expires and expires < now:
            removed += 1
            continue
        kept.append(item)
    return kept, removed


def allocation_sort_key(item: dict[str, Any]) -> datetime:
    last_seen = parse_iso(str(item.get("last_seen_at") or ""))
    allocated = parse_iso(str(item.get("allocated_at") or ""))
    return last_seen or allocated or datetime.fromtimestamp(0, tz=timezone.utc)


def output(payload: dict[str, Any]) -> None:
    print(json.dumps(payload, ensure_ascii=False, indent=2))


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Allocate/reuse/release branch preview slots for issue-driven workflows."
    )
    parser.add_argument("--state-file", default=".tmp/preview-slots.json")
    parser.add_argument("--slots", required=True, help="Comma-separated slot names, e.g. preview-a,preview-b")
    parser.add_argument("--repo", default="", help="owner/repo")
    parser.add_argument("--issue", type=int, default=None, help="Issue number")
    parser.add_argument("--branch", default="", help="Branch name")
    parser.add_argument("--slot", default="", help="Slot name (optional release filter)")
    parser.add_argument("--ttl-hours", type=int, default=24, help="Allocation TTL in hours")
    parser.add_argument("--url-template", default="", help="Optional URL template, use {slot} placeholder")
    parser.add_argument("--release", action="store_true", help="Release matching allocation(s)")
    parser.add_argument("--list", action="store_true", help="List active allocations and free slots")
    parser.add_argument("--evict-oldest", action="store_true", help="Evict oldest allocation when slots are full")
    args = parser.parse_args()

    state_path = Path(args.state_file)
    slots = parse_slots(args.slots)
    state = load_state(state_path)

    now = utc_now()
    allocations_raw = [item for item in state.get("allocations", []) if isinstance(item, dict)]
    allocations, pruned_count = prune_expired(allocations_raw, now)
    state["allocations"] = allocations

    if args.list:
        used_slots = {str(item.get("slot") or "") for item in allocations}
        free_slots = [slot for slot in slots if slot not in used_slots]
        output(
            {
                "action": "list",
                "state_file": str(state_path.as_posix()),
                "pruned_expired": pruned_count,
                "total_active": len(allocations),
                "active_allocations": allocations,
                "free_slots": free_slots,
            }
        )
        save_state(state_path, state)
        return

    if args.release:
        target_repo = args.repo.strip()
        target_branch = args.branch.strip()
        target_slot = args.slot.strip()
        target_issue = args.issue

        kept: list[dict[str, Any]] = []
        released: list[dict[str, Any]] = []
        for item in allocations:
            match = True
            if target_repo:
                match = match and str(item.get("repo") or "") == target_repo
            if target_branch:
                match = match and str(item.get("branch") or "") == target_branch
            if target_slot:
                match = match and str(item.get("slot") or "") == target_slot
            if target_issue is not None:
                match = match and int(item.get("issue") or -1) == target_issue

            if match:
                released.append(item)
            else:
                kept.append(item)

        state["allocations"] = kept
        save_state(state_path, state)
        output(
            {
                "action": "release",
                "state_file": str(state_path.as_posix()),
                "released_count": len(released),
                "released": released,
                "remaining_count": len(kept),
            }
        )
        return

    repo = args.repo.strip()
    branch = args.branch.strip()
    if not repo or not branch:
        print("allocate requires --repo and --branch", file=sys.stderr)
        sys.exit(2)

    # Reuse existing allocation for same repo+branch.
    existing = next(
        (item for item in allocations if str(item.get("repo") or "") == repo and str(item.get("branch") or "") == branch),
        None,
    )
    if existing is not None:
        existing["issue"] = args.issue if args.issue is not None else existing.get("issue")
        existing["last_seen_at"] = to_iso(now)
        existing["expires_at"] = to_iso(now + timedelta(hours=max(1, args.ttl_hours)))
        if args.url_template:
            existing["url"] = render_url(str(existing.get("slot") or ""), args.url_template)
        save_state(state_path, state)
        output(
            {
                "action": "allocate",
                "reused": True,
                "state_file": str(state_path.as_posix()),
                "allocation": existing,
                "pruned_expired": pruned_count,
            }
        )
        return

    used_slots = {str(item.get("slot") or "") for item in allocations}
    free_slots = [slot for slot in slots if slot not in used_slots]
    evicted: dict[str, Any] | None = None

    if not free_slots:
        if not args.evict_oldest:
            output(
                {
                    "action": "allocate",
                    "reused": False,
                    "allocated": False,
                    "reason": "no_free_slots",
                    "state_file": str(state_path.as_posix()),
                    "active_allocations": allocations,
                }
            )
            sys.exit(3)
        oldest = sorted(allocations, key=allocation_sort_key)[0]
        evicted = oldest
        allocations.remove(oldest)
        free_slots = [str(oldest.get("slot") or "")]

    slot = free_slots[0]
    allocation = {
        "repo": repo,
        "issue": args.issue,
        "branch": branch,
        "slot": slot,
        "env_id": slot,
        "url": render_url(slot, args.url_template),
        "allocated_at": to_iso(now),
        "last_seen_at": to_iso(now),
        "expires_at": to_iso(now + timedelta(hours=max(1, args.ttl_hours))),
        "status": "active",
    }
    allocations.append(allocation)
    state["allocations"] = allocations
    save_state(state_path, state)
    output(
        {
            "action": "allocate",
            "reused": False,
            "allocated": True,
            "state_file": str(state_path.as_posix()),
            "allocation": allocation,
            "evicted": evicted,
            "pruned_expired": pruned_count,
        }
    )


if __name__ == "__main__":
    main()