"""
Agent ERP (Enterprise Resource Planning) for Brynq.

Tracks resources and capacity across AI agents:
- Token budgets: estimated token usage per agent per task
- Workload tracking: active tasks, estimated completion times
- Capacity planning: available capacity per agent (max concurrent vs current)
- Cost tracking: estimated cost per task, per agent, per project
- Resource allocation: assign resource limits to agents
- Utilization reports: who's overloaded, who's idle

Storage:
  state/broker/erp/agents/{session_id}.json — agent resource profile
  state/broker/erp/ledger.jsonl — all resource events (allocations, completions, costs)
"""

import json
import os
import uuid
from datetime import datetime, timezone, timedelta
from pathlib import Path
from typing import Optional

from brynq.server import (
    BROKER_DIR,
    CHANNELS_DIR,
    SESSION_ID,
    SESSION_NAME,
    PLATFORM,
    _append_jsonl_sync,
    _write_json_sync,
    _read_jsonl_sync,
)

# ── Paths ──────────────────────────────────────────────────────────────────

ERP_DIR = BROKER_DIR / "erp"
AGENTS_DIR = ERP_DIR / "agents"
LEDGER_FILE = ERP_DIR / "ledger.jsonl"

for d in (ERP_DIR, AGENTS_DIR):
    d.mkdir(parents=True, exist_ok=True)

# ── Default cost rates (USD per 1K tokens) ────────────────────────────────

DEFAULT_INPUT_COST_PER_1K = 0.003
DEFAULT_OUTPUT_COST_PER_1K = 0.015


# ── Internal Helpers ──────────────────────────────────────────────────────


def _load_agent_profile(session_id: str) -> Optional[dict]:
    """Load an agent's resource profile from disk."""
    path = AGENTS_DIR / f"{session_id}.json"
    if not path.exists():
        return None
    try:
        return json.loads(path.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return None


def _save_agent_profile(profile: dict) -> None:
    """Save an agent's resource profile to disk."""
    sid = profile["session_id"]
    _write_json_sync(AGENTS_DIR / f"{sid}.json", profile)


def _ensure_agent_profile(session_id: str) -> dict:
    """Load or create a default agent resource profile."""
    profile = _load_agent_profile(session_id)
    if profile is not None:
        return profile

    now = datetime.now(timezone.utc).isoformat()
    profile = {
        "session_id": session_id,
        "max_concurrent_tasks": 5,
        "token_budget_per_hour": 100_000,
        "active_tasks": {},
        "total_tasks_completed": 0,
        "total_tokens_used": 0,
        "total_cost_usd": 0.0,
        "created": now,
        "updated": now,
    }
    _save_agent_profile(profile)
    return profile


def _log_event(event_type: str, data: dict) -> None:
    """Append a resource event to the ledger."""
    entry = {
        "id": uuid.uuid4().hex[:12],
        "timestamp": datetime.now(timezone.utc).isoformat(),
        "event_type": event_type,
        "session_id": SESSION_ID,
        "session_name": SESSION_NAME,
        "platform": PLATFORM,
        **data,
    }
    _append_jsonl_sync(LEDGER_FILE, entry)


def _estimate_cost(tokens: int, split_ratio: float = 0.5) -> float:
    """Estimate USD cost for a given token count.

    Assumes *split_ratio* of tokens are input and the rest are output.
    Uses conservative default rates.
    """
    input_tokens = tokens * split_ratio
    output_tokens = tokens * (1.0 - split_ratio)
    cost = (input_tokens / 1000) * DEFAULT_INPUT_COST_PER_1K
    cost += (output_tokens / 1000) * DEFAULT_OUTPUT_COST_PER_1K
    return round(cost, 6)


# ── Public API ────────────────────────────────────────────────────────────


def set_capacity(
    max_concurrent_tasks: int = 5,
    token_budget_per_hour: int = 100_000,
    session_id: Optional[str] = None,
) -> dict:
    """Set an agent's resource limits.

    Args:
        max_concurrent_tasks: Maximum number of tasks the agent can handle
            concurrently.
        token_budget_per_hour: Maximum estimated token throughput per hour.
        session_id: Target agent. Defaults to current session.

    Returns:
        The updated agent resource profile.
    """
    sid = session_id or SESSION_ID
    profile = _ensure_agent_profile(sid)

    profile["max_concurrent_tasks"] = max_concurrent_tasks
    profile["token_budget_per_hour"] = token_budget_per_hour
    profile["updated"] = datetime.now(timezone.utc).isoformat()
    _save_agent_profile(profile)

    _log_event("capacity_set", {
        "target_session_id": sid,
        "max_concurrent_tasks": max_concurrent_tasks,
        "token_budget_per_hour": token_budget_per_hour,
    })

    return profile


def get_capacity(session_id: Optional[str] = None) -> dict:
    """Get an agent's current capacity and utilization.

    Args:
        session_id: Agent to query. Defaults to current session.

    Returns:
        Dict with capacity limits, active task count, utilization percentage,
        and remaining budget.
    """
    sid = session_id or SESSION_ID
    profile = _ensure_agent_profile(sid)

    active_count = len(profile.get("active_tasks", {}))
    max_tasks = profile.get("max_concurrent_tasks", 5)
    token_budget = profile.get("token_budget_per_hour", 100_000)

    # Sum estimated tokens for active tasks
    active_tokens = sum(
        t.get("estimated_tokens", 0)
        for t in profile.get("active_tasks", {}).values()
    )

    utilization_pct = round((active_count / max(max_tasks, 1)) * 100, 1)

    return {
        "session_id": sid,
        "max_concurrent_tasks": max_tasks,
        "active_task_count": active_count,
        "available_slots": max(max_tasks - active_count, 0),
        "utilization_pct": utilization_pct,
        "token_budget_per_hour": token_budget,
        "estimated_active_tokens": active_tokens,
        "remaining_token_budget": max(token_budget - active_tokens, 0),
        "total_tasks_completed": profile.get("total_tasks_completed", 0),
        "total_tokens_used": profile.get("total_tokens_used", 0),
        "total_cost_usd": profile.get("total_cost_usd", 0.0),
    }


def allocate_task(
    task_id: str,
    estimated_tokens: int = 0,
    estimated_hours: float = 0.0,
    project: str = "",
    session_id: Optional[str] = None,
) -> dict:
    """Allocate resources for a task to an agent.

    Records the task in the agent's active task set and logs the allocation.

    Args:
        task_id: Unique identifier of the task being allocated.
        estimated_tokens: Estimated token usage for the task.
        estimated_hours: Estimated hours to complete.
        project: Optional project name for cost grouping.
        session_id: Target agent. Defaults to current session.

    Returns:
        Dict with allocation details including whether the agent is now
        at capacity.
    """
    sid = session_id or SESSION_ID
    profile = _ensure_agent_profile(sid)

    active = profile.get("active_tasks", {})
    if task_id in active:
        return {"error": f"Task {task_id} already allocated to {sid}"}

    max_tasks = profile.get("max_concurrent_tasks", 5)
    if len(active) >= max_tasks:
        return {
            "error": f"Agent {sid} at capacity ({len(active)}/{max_tasks} tasks)",
            "active_task_count": len(active),
            "max_concurrent_tasks": max_tasks,
        }

    now = datetime.now(timezone.utc)
    estimated_completion = None
    if estimated_hours > 0:
        estimated_completion = (now + timedelta(hours=estimated_hours)).isoformat()

    estimated_cost = _estimate_cost(estimated_tokens) if estimated_tokens > 0 else 0.0

    task_entry = {
        "task_id": task_id,
        "estimated_tokens": estimated_tokens,
        "estimated_hours": estimated_hours,
        "estimated_cost_usd": estimated_cost,
        "project": project,
        "allocated_at": now.isoformat(),
        "estimated_completion": estimated_completion,
    }

    active[task_id] = task_entry
    profile["active_tasks"] = active
    profile["updated"] = now.isoformat()
    _save_agent_profile(profile)

    _log_event("task_allocated", {
        "target_session_id": sid,
        "task_id": task_id,
        "estimated_tokens": estimated_tokens,
        "estimated_hours": estimated_hours,
        "estimated_cost_usd": estimated_cost,
        "project": project,
    })

    return {
        "status": "allocated",
        "task_id": task_id,
        "session_id": sid,
        "active_task_count": len(active),
        "max_concurrent_tasks": max_tasks,
        "at_capacity": len(active) >= max_tasks,
        "estimated_cost_usd": estimated_cost,
        "estimated_completion": estimated_completion,
    }


def complete_task(
    task_id: str,
    actual_tokens: int = 0,
    actual_hours: float = 0.0,
    session_id: Optional[str] = None,
) -> dict:
    """Record actual resource usage when a task completes.

    Removes the task from the agent's active set and updates lifetime totals.

    Args:
        task_id: The task that completed.
        actual_tokens: Actual token count used.
        actual_hours: Actual hours spent.
        session_id: Agent that completed the task. Defaults to current session.

    Returns:
        Dict with completion details, cost, and variance from estimates.
    """
    sid = session_id or SESSION_ID
    profile = _ensure_agent_profile(sid)

    active = profile.get("active_tasks", {})
    task_entry = active.pop(task_id, None)
    if task_entry is None:
        return {"error": f"Task {task_id} not found in active tasks for {sid}"}

    now = datetime.now(timezone.utc).isoformat()
    actual_cost = _estimate_cost(actual_tokens) if actual_tokens > 0 else 0.0

    # Calculate variance from estimates
    est_tokens = task_entry.get("estimated_tokens", 0)
    est_hours = task_entry.get("estimated_hours", 0.0)
    est_cost = task_entry.get("estimated_cost_usd", 0.0)

    token_variance = actual_tokens - est_tokens if est_tokens > 0 else 0
    hours_variance = round(actual_hours - est_hours, 4) if est_hours > 0 else 0.0
    cost_variance = round(actual_cost - est_cost, 6) if est_cost > 0 else 0.0

    # Update lifetime totals
    profile["active_tasks"] = active
    profile["total_tasks_completed"] = profile.get("total_tasks_completed", 0) + 1
    profile["total_tokens_used"] = profile.get("total_tokens_used", 0) + actual_tokens
    profile["total_cost_usd"] = round(
        profile.get("total_cost_usd", 0.0) + actual_cost, 6
    )
    profile["updated"] = now
    _save_agent_profile(profile)

    _log_event("task_completed", {
        "target_session_id": sid,
        "task_id": task_id,
        "actual_tokens": actual_tokens,
        "actual_hours": actual_hours,
        "actual_cost_usd": actual_cost,
        "estimated_tokens": est_tokens,
        "estimated_hours": est_hours,
        "estimated_cost_usd": est_cost,
        "token_variance": token_variance,
        "hours_variance": hours_variance,
        "cost_variance": cost_variance,
        "project": task_entry.get("project", ""),
    })

    return {
        "status": "completed",
        "task_id": task_id,
        "session_id": sid,
        "actual_tokens": actual_tokens,
        "actual_hours": actual_hours,
        "actual_cost_usd": actual_cost,
        "token_variance": token_variance,
        "hours_variance": hours_variance,
        "cost_variance": cost_variance,
        "active_task_count": len(active),
        "total_tasks_completed": profile["total_tasks_completed"],
    }


def get_utilization() -> list[dict]:
    """List all agents with their utilization percentage.

    Scans all agent profiles in the ERP directory and returns a sorted list
    from most utilized to least.

    Returns:
        List of dicts with session_id, active tasks, max tasks,
        utilization percentage, and a status label.
    """
    agents = []
    for f in sorted(AGENTS_DIR.glob("*.json")):
        try:
            profile = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        sid = profile.get("session_id", f.stem)
        active_count = len(profile.get("active_tasks", {}))
        max_tasks = profile.get("max_concurrent_tasks", 5)
        utilization = round((active_count / max(max_tasks, 1)) * 100, 1)

        if utilization >= 90:
            status = "overloaded"
        elif utilization >= 60:
            status = "busy"
        elif utilization > 0:
            status = "active"
        else:
            status = "idle"

        agents.append({
            "session_id": sid,
            "active_task_count": active_count,
            "max_concurrent_tasks": max_tasks,
            "utilization_pct": utilization,
            "status": status,
            "token_budget_per_hour": profile.get("token_budget_per_hour", 0),
            "total_tasks_completed": profile.get("total_tasks_completed", 0),
            "total_cost_usd": profile.get("total_cost_usd", 0.0),
        })

    agents.sort(key=lambda a: a["utilization_pct"], reverse=True)
    return agents


def get_cost_report(period_hours: float = 24) -> dict:
    """Generate a cost summary over a time period.

    Reads the ledger and aggregates completed-task costs by agent and project.

    Args:
        period_hours: How far back to look in hours. Defaults to 24.

    Returns:
        Dict with total cost, per-agent breakdown, per-project breakdown,
        task count, and token totals.
    """
    entries = _read_jsonl_sync(LEDGER_FILE)
    cutoff = (
        datetime.now(timezone.utc) - timedelta(hours=period_hours)
    ).isoformat()

    by_agent: dict[str, dict] = {}
    by_project: dict[str, dict] = {}
    total_cost = 0.0
    total_tokens = 0
    task_count = 0

    for entry in entries:
        if entry.get("event_type") != "task_completed":
            continue
        ts = entry.get("timestamp", "")
        if ts < cutoff:
            continue

        cost = entry.get("actual_cost_usd", 0.0)
        tokens = entry.get("actual_tokens", 0)
        agent = entry.get("target_session_id", "unknown")
        project = entry.get("project", "") or "unassigned"

        total_cost += cost
        total_tokens += tokens
        task_count += 1

        # Per-agent accumulation
        if agent not in by_agent:
            by_agent[agent] = {"cost_usd": 0.0, "tokens": 0, "tasks": 0}
        by_agent[agent]["cost_usd"] = round(by_agent[agent]["cost_usd"] + cost, 6)
        by_agent[agent]["tokens"] += tokens
        by_agent[agent]["tasks"] += 1

        # Per-project accumulation
        if project not in by_project:
            by_project[project] = {"cost_usd": 0.0, "tokens": 0, "tasks": 0}
        by_project[project]["cost_usd"] = round(
            by_project[project]["cost_usd"] + cost, 6
        )
        by_project[project]["tokens"] += tokens
        by_project[project]["tasks"] += 1

    return {
        "period_hours": period_hours,
        "total_cost_usd": round(total_cost, 6),
        "total_tokens": total_tokens,
        "task_count": task_count,
        "by_agent": by_agent,
        "by_project": by_project,
    }


def find_available_agents(min_capacity: int = 1) -> list[dict]:
    """Find agents with spare capacity for new tasks.

    Args:
        min_capacity: Minimum number of free task slots required.

    Returns:
        List of agents with at least *min_capacity* open slots, sorted by
        most available capacity first.
    """
    available = []
    for f in sorted(AGENTS_DIR.glob("*.json")):
        try:
            profile = json.loads(f.read_text("utf-8"))
        except (json.JSONDecodeError, OSError):
            continue

        sid = profile.get("session_id", f.stem)
        active_count = len(profile.get("active_tasks", {}))
        max_tasks = profile.get("max_concurrent_tasks", 5)
        free_slots = max(max_tasks - active_count, 0)

        if free_slots < min_capacity:
            continue

        utilization = round((active_count / max(max_tasks, 1)) * 100, 1)

        available.append({
            "session_id": sid,
            "available_slots": free_slots,
            "active_task_count": active_count,
            "max_concurrent_tasks": max_tasks,
            "utilization_pct": utilization,
            "token_budget_per_hour": profile.get("token_budget_per_hour", 0),
            "total_tasks_completed": profile.get("total_tasks_completed", 0),
        })

    # Most available capacity first
    available.sort(key=lambda a: a["available_slots"], reverse=True)
    return available
