#!/usr/bin/env python3
# Copyright (c) 2026 PULSE Contributors. Proprietary. All rights reserved.
"""
Causal Crew Composer — Brain V8 Phase 3 (D3)

Evidence-based swarm composition. Queries knowledge graph for causal risks
and scores crew using Tier Alignment + provider-aware limits.
"""
from __future__ import annotations

import logging

from pulse.core.brain_utils import load_json_dict, TASK_BOARD_FILE

log = logging.getLogger("pulse.crew_composer")

# --- Tier / Complexity Mapping ---
_COMPLEXITY_TIER = {"low": "fast", "medium": "standard", "high": "premium"}
_TIER_RANK = {"fast": 0, "standard": 1, "premium": 2}


def _tier_covers(crew_tier: str, task_tier: str) -> bool:
    """True if crew_tier is >= task_tier in capability ranking."""
    return _TIER_RANK.get(crew_tier, 0) >= _TIER_RANK.get(task_tier, 0)


def get_pending_tasks() -> list[dict]:
    """Read task board and return open tasks with domain/complexity."""
    board = load_json_dict(TASK_BOARD_FILE)
    tasks = []
    for dispatch in board.get("dispatches", []):
        for t in dispatch.get("tasks", []):
            status = t.get("status", "")
            if status in ("open", "escalated"):
                tasks.append({
                    "task_id": t.get("task_id") or t.get("id", ""),
                    "title": t.get("title", ""),
                    "domain": t.get("domain", "general"),
                    "complexity": t.get("complexity", "low"),
                    "recommended_tier": t.get("recommended_tier", "fast"),
                })
    return tasks


def _query_graph_risks(domains: set[str]) -> list[dict]:
    """Query knowledge graph for CAUSES/LEADS_TO risks related to task domains."""
    try:
        from pulse.graph.knowledge_graph_query import query_causes, query_prevents
    except ImportError:
        return []

    risks = []
    for domain in domains:
        causes = query_causes(f"failure in {domain}", max_results=3)
        for c in causes:
            if c.get("confidence", 0) >= 0.4:
                risks.append({
                    "domain": domain,
                    "type": "cause",
                    "text": c["text"],
                    "relationship": c.get("relationship", "CAUSES"),
                    "confidence": c.get("confidence", 0.5),
                })
        prevents = query_prevents(f"failure in {domain}", max_results=3)
        for p in prevents:
            if p.get("confidence", 0) >= 0.4:
                risks.append({
                    "domain": domain,
                    "type": "protection",
                    "text": p["text"],
                    "relationship": "PREVENTS",
                    "confidence": p.get("confidence", 0.5),
                })
    return risks


def _check_agent_weaknesses(crew: list[dict], domains: set[str]) -> list[dict]:
    """Check D2 agent profiles for domain weaknesses in current crew."""
    try:
        from pulse.brain.agent_procedures import get_weak_domains
    except ImportError:
        return []

    warnings = []
    for member in crew:
        model = member.get("model", "")
        weak = get_weak_domains(model)
        for w in weak:
            if w["domain"] in domains:
                warnings.append({
                    "agent": model,
                    "domain": w["domain"],
                    "success_rate": w["success_rate"],
                    "total": w["total"],
                })
    return warnings


def _score_tier_alignment(crew: list[dict], tasks: list[dict]) -> tuple[float, list[str]]:
    """Score Tier Alignment: high tasks need premium, low tasks reward fast. Returns (score, warnings)."""
    if not crew or not tasks:
        return 0.5, []
    tier_needs = {"fast": 0, "standard": 0, "premium": 0}
    for t in tasks:
        tier_needs[_COMPLEXITY_TIER.get(t.get("complexity", "low"), "fast")] += 1
    crew_tiers = {"fast": 0, "standard": 0, "premium": 0}
    for m in crew:
        crew_tiers[m.get("tier", "fast")] = crew_tiers.get(m.get("tier", "fast"), 0) + m.get("count", 1)
    warnings = []
    score = 0.5

    # Rule 1: High-complexity tasks NEED premium agents
    if tier_needs["premium"] > 0 and crew_tiers.get("premium", 0) == 0:
        score -= 0.3
        warnings.append(
            f"{tier_needs['premium']} high-complexity task(s) but NO premium agents in crew"
        )
    elif tier_needs["premium"] > 0 and crew_tiers.get("premium", 0) > 0:
        score += 0.2

    # Rule 2: Medium-complexity tasks NEED standard+ agents
    needs_standard = tier_needs["standard"]
    has_standard_plus = crew_tiers.get("standard", 0) + crew_tiers.get("premium", 0)
    if needs_standard > 0 and has_standard_plus == 0:
        score -= 0.2
        warnings.append(
            f"{needs_standard} standard-complexity task(s) but only fast-tier agents"
        )
    elif needs_standard > 0 and has_standard_plus > 0:
        score += 0.15

    # Rule 3: All-low tasks with all-fast crew = OPTIMAL (cost efficiency)
    if tier_needs["premium"] == 0 and tier_needs["standard"] == 0:
        total_crew = sum(crew_tiers.values())
        if crew_tiers.get("fast", 0) == total_crew and total_crew > 0:
            score += 0.3  # Perfect cost alignment
        elif crew_tiers.get("premium", 0) > 0:
            score -= 0.1
            warnings.append(
                "All tasks are low-complexity — premium agents waste tokens"
            )

    return max(0.0, min(1.0, score)), warnings


OLLAMA_FREE_LIMIT = 5  # Max concurrent agents on Ollama free plan


def _check_provider_limits(crew: list[dict]) -> tuple[float, list[str], bool]:
    """Check Ollama free tier cap. Returns (penalty, warnings, is_ollama_only)."""
    warnings = []
    penalty = 0.0
    ollama_count = sum(
        m.get("count", 1) for m in crew if m.get("provider", "") == "ollama"
    )
    if ollama_count > OLLAMA_FREE_LIMIT:
        penalty = 0.15
        warnings.append(
            f"Ollama free plan allows max {OLLAMA_FREE_LIMIT} concurrent — "
            f"crew has {ollama_count} Ollama agents"
        )
    has_paid = any(m.get("provider", "") in ("claude", "gemini", "openai") for m in crew)
    is_ollama_only = not has_paid and ollama_count > 0
    return penalty, warnings, is_ollama_only


def score_crew(crew: list[dict], tasks: list[dict] | None = None) -> dict:
    """Score crew via Tier Alignment + graph risks + provider limits. Returns analysis dict."""
    if tasks is None:
        tasks = get_pending_tasks()

    if not crew:
        return {"score": 0.0, "warnings": ["No agents in crew"], "suggestions": [], "task_summary": {}}

    domains = {t.get("domain", "general") for t in tasks}
    warnings = []
    suggestions = []

    # 1. Tier Alignment scoring (no diversity bonus — Gemini fix)
    tier_score, tier_warnings = _score_tier_alignment(crew, tasks)
    warnings.extend(tier_warnings)

    # 2. Provider limits (Ollama free tier cap)
    provider_penalty, provider_warnings, is_ollama_only = _check_provider_limits(crew)
    warnings.extend(provider_warnings)

    # 3. Knowledge Graph risk analysis
    graph_risks = _query_graph_risks(domains)
    cause_risks = [r for r in graph_risks if r["type"] == "cause"]
    protections = [r for r in graph_risks if r["type"] == "protection"]

    for risk in cause_risks[:3]:
        warnings.append(
            f"{risk['domain']}: {risk['text']} ({risk['relationship']}, "
            f"confidence: {risk['confidence']:.0%})"
        )
    risk_penalty = len(cause_risks) * 0.05
    protection_bonus = min(len(protections) * 0.03, 0.1)

    # 4. D2 agent weakness check
    agent_warnings = _check_agent_weaknesses(crew, domains)
    for aw in agent_warnings[:2]:
        warnings.append(
            f"Agent {aw['agent']} has {aw['success_rate']:.0%} success on {aw['domain']} "
            f"({aw['total']} tasks)"
        )
    weakness_penalty = len(agent_warnings) * 0.05

    # 5. Generate provider-aware suggestions
    if any("high-complexity" in w.lower() for w in tier_warnings):
        if is_ollama_only:
            suggestions.append(
                "High-complexity tasks need premium tier — Ollama cloud models "
                "(qwen3-coder-next, kimi-k2.5) are free premium options"
            )
        else:
            suggestions.append("Add 1x premium-tier agent for high-complexity oversight")
    if any("only fast" in w.lower() for w in tier_warnings):
        if is_ollama_only:
            suggestions.append("Swap 1x fast agent for an Ollama cloud model (standard+)")
        else:
            suggestions.append("Add 1x standard-tier agent for medium-complexity tasks")
    if any("waste tokens" in w.lower() for w in tier_warnings):
        suggestions.append("Replace premium agents with fast-tier for cost savings")
    if any("ollama free plan" in w.lower() for w in provider_warnings):
        suggestions.append(
            f"Reduce Ollama agents to {OLLAMA_FREE_LIMIT} or upgrade Ollama plan"
        )
    for aw in agent_warnings[:1]:
        suggestions.append(
            f"Consider replacing {aw['agent']} — weak on {aw['domain']} tasks"
        )

    # Final score
    final_score = (tier_score - risk_penalty + protection_bonus
                   - weakness_penalty - provider_penalty)
    final_score = max(0.0, min(1.0, final_score))

    # Task summary
    complexity_counts = {}
    domain_counts = {}
    for t in tasks:
        c = t.get("complexity", "low")
        d = t.get("domain", "general")
        complexity_counts[c] = complexity_counts.get(c, 0) + 1
        domain_counts[d] = domain_counts.get(d, 0) + 1

    return {
        "score": round(final_score, 2),
        "warnings": warnings,
        "suggestions": suggestions,
        "task_summary": {
            "total": len(tasks),
            "by_complexity": complexity_counts,
            "by_domain": domain_counts,
        },
        "graph_risks": len(cause_risks),
        "graph_protections": len(protections),
        "ollama_only": is_ollama_only,
    }


def format_crew_analysis(analysis: dict) -> str:
    """Format crew analysis for CLI display."""
    score = analysis.get("score", 0.5)
    warnings = analysis.get("warnings", [])
    suggestions = analysis.get("suggestions", [])
    summary = analysis.get("task_summary", {})

    if score >= 0.7:
        label = "LOW RISK"
    elif score >= 0.4:
        label = "MODERATE RISK"
    else:
        label = "HIGH RISK"

    lines = [
        "",
        "  --- CREW ANALYSIS ---",
        f"  Score: {score:.2f}/1.0 ({label})",
    ]

    total = summary.get("total", 0)
    if total > 0:
        domain_parts = [f"{c} {d}" for d, c in summary.get("by_domain", {}).items()]
        lines.append(f"  Pending tasks: {total} ({', '.join(domain_parts)})")

    if warnings:
        lines.append("")
        for w in warnings[:4]:
            lines.append(f"  !  {w}")

    if suggestions:
        lines.append("")
        for s in suggestions[:3]:
            lines.append(f"  +  {s}")

    lines.append("  " + "-" * 40)
    return "\n".join(lines)
