"""
CrewAI Service - Canonical Multi-Role Orchestration v2.0
Variant A: Profiles per top-level agent
"""
import os
import json
import time
import asyncio
import logging
import httpx
from typing import Dict, Any, List, Optional
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

# Setup logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# Import registry loader (v2 with profiles)
from registry_loader import (
    load_teams_config,
    get_team_members, get_synthesis_config, get_delegation_config,
    get_team_settings, can_delegate_to, is_orchestrator,
    get_agent_profiles, get_default_profile, select_profile,
    get_profile_config, get_all_agents_summary
)

app = FastAPI(title="CrewAI Service", version="2.0.0")

# Configuration
ROUTER_URL = os.getenv("ROUTER_URL", "http://dagi-staging-router:8000")
DEFAULT_MAX_CONCURRENCY = int(os.getenv("MAX_CONCURRENT_ROLES", "3"))
LLM_TIMEOUT = int(os.getenv("LLM_TIMEOUT", "120"))


# Request/Response models
class CrewRunRequest(BaseModel):
    orchestrator_id: str
    task: str
    context: Dict[str, Any] = {}
    profile: Optional[str] = None  # NEW: explicit profile selection

class CrewRunResponse(BaseModel):
    success: bool
    result: Optional[str] = None
    error: Optional[str] = None
    meta: Dict[str, Any] = {}

class DelegationRequest(BaseModel):
    orchestrator_id: str
    target_agent_id: str
    task: str
    context: Dict[str, Any] = {}
    hops_remaining: int = 2


async def call_internal_llm(
    prompt: str,
    system_prompt: str = None,
    role_context: str = None,
    llm_profile: str = "reasoning"
) -> str:
    """Call Router internal LLM endpoint for a single role"""
    url = f"{ROUTER_URL}/internal/llm/complete"
    
    payload = {
        "prompt": prompt,
        "llm_profile": llm_profile,
        "max_tokens": 2048,
        "temperature": 0.3
    }
    if system_prompt:
        payload["system_prompt"] = system_prompt
    if role_context:
        payload["role_context"] = role_context
    
    async with httpx.AsyncClient(timeout=LLM_TIMEOUT) as client:
        try:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()
            return data.get("text", "")
        except Exception as e:
            logger.error(f"Internal LLM call failed: {e}")
            raise


async def delegate_to_agent(
    orchestrator_id: str,
    target_agent_id: str,
    task: str,
    context: Dict[str, Any] = {},
    hops_remaining: int = 2,
    profile: str = None
) -> Optional[str]:
    """Delegate task to another top-level agent via Router"""
    if not can_delegate_to(orchestrator_id, target_agent_id, profile):
        logger.warning(f"Delegation not allowed: {orchestrator_id} -> {target_agent_id}")
        return None
    
    if hops_remaining <= 0:
        logger.warning(f"Max delegation hops reached for {orchestrator_id}")
        return None
    
    url = f"{ROUTER_URL}/v1/agents/{target_agent_id}/infer"
    
    delegation_cfg = get_delegation_config(orchestrator_id, profile) or {}
    attach_headers = delegation_cfg.get("attach_headers", {})
    
    payload = {
        "prompt": task,
        "metadata": {
            "handoff_from": attach_headers.get("handoff_from", orchestrator_id),
            "hops_remaining": hops_remaining - 1,
        }
    }
    
    logger.info(f"DELEGATION: {orchestrator_id} -> {target_agent_id} (hops={hops_remaining})")
    
    async with httpx.AsyncClient(timeout=180) as client:
        try:
            resp = await client.post(url, json=payload)
            resp.raise_for_status()
            data = resp.json()
            return data.get("response", "")
        except Exception as e:
            logger.error(f"Delegation to {target_agent_id} failed: {e}")
            return None


async def execute_role(
    role_config: Dict[str, Any],
    task: str,
    context: Dict[str, Any],
    semaphore: asyncio.Semaphore
) -> Dict[str, Any]:
    """Execute a single role with rate limiting"""
    role_id = role_config.get("id", "unknown")
    role_context = role_config.get("role_context", role_id)
    llm_profile = role_config.get("llm_profile", "reasoning")
    system_prompt = role_config.get("system_prompt", "")
    
    memory_brief = context.get("memory_brief", {})
    memory_str = json.dumps(memory_brief, ensure_ascii=False)[:500] if memory_brief else ""
    
    prompt = f"Task: {task}\n\nContext: {memory_str}\n\nYour role: {role_context}\n\nProvide your analysis and recommendations."

    async with semaphore:
        t0 = time.time()
        try:
            logger.info(f"ROLE START: {role_context} (profile={llm_profile})")
            result = await call_internal_llm(
                prompt=prompt,
                system_prompt=system_prompt,
                role_context=role_context,
                llm_profile=llm_profile
            )
            elapsed = time.time() - t0
            logger.info(f"ROLE DONE: {role_context} ({elapsed:.1f}s)")
            return {
                "role_id": role_id,
                "role_context": role_context,
                "result": result,
                "elapsed_seconds": elapsed,
                "success": True
            }
        except Exception as e:
            elapsed = time.time() - t0
            logger.error(f"ROLE ERROR: {role_context}: {e}")
            return {
                "role_id": role_id,
                "role_context": role_context,
                "result": None,
                "error": str(e),
                "elapsed_seconds": elapsed,
                "success": False
            }


async def run_crew_canonical(
    orchestrator_id: str,
    task: str,
    context: Dict[str, Any],
    profile: str = None
) -> CrewRunResponse:
    """Execute multi-role orchestration for an agent with profile selection"""
    
    # Select profile (auto or explicit)
    if profile is None:
        profile = select_profile(orchestrator_id, task)
    
    logger.info(f"CREW RUN: {orchestrator_id} profile={profile}")
    
    # Get team config for selected profile
    team_members = get_team_members(orchestrator_id, profile)
    synthesis_config = get_synthesis_config(orchestrator_id, profile)
    team_settings = get_team_settings(orchestrator_id, profile)
    
    if not team_members:
        return CrewRunResponse(
            success=False,
            error=f"No team members for {orchestrator_id}.{profile}",
            meta={"orchestrator_id": orchestrator_id, "profile": profile}
        )
    
    team_name = team_settings.get("team_name", f"{orchestrator_id} team")
    parallel_roles = team_settings.get("parallel_roles", True)
    max_concurrency = team_settings.get("max_concurrency", DEFAULT_MAX_CONCURRENCY)
    
    logger.info(f"Team: {team_name}, Roles: {len(team_members)}, Parallel: {parallel_roles}, MaxConc: {max_concurrency}")
    
    t0 = time.time()
    role_results = []
    
    if parallel_roles:
        # Parallel execution with semaphore
        semaphore = asyncio.Semaphore(max_concurrency)
        tasks = [
            execute_role(member, task, context, semaphore)
            for member in team_members
        ]
        role_results = await asyncio.gather(*tasks, return_exceptions=True)
        role_results = [r if isinstance(r, dict) else {"error": str(r), "success": False} for r in role_results]
    else:
        # Sequential execution
        semaphore = asyncio.Semaphore(1)
        for member in team_members:
            res = await execute_role(member, task, context, semaphore)
            role_results.append(res)
    
    # Build synthesis context
    successful_results = [r for r in role_results if r.get("success")]
    failed_results = [r for r in role_results if not r.get("success")]
    
    synthesis_context = "\n\n".join([
        f"## {r['role_context']}\n{r['result']}"
        for r in successful_results
    ])
    
    # Synthesis
    synthesis_prompt = synthesis_config.get("system_prompt", "")
    synthesis_role = synthesis_config.get("role_context", "Synthesis")
    synthesis_llm = synthesis_config.get("llm_profile", "reasoning")
    
    final_prompt = f"""Task: {task}

Team Analysis:
{synthesis_context}

Synthesize the above into a coherent, actionable response. Include:
- Key findings
- Recommendations
- Risks/limitations
- Next steps"""

    try:
        final_result = await call_internal_llm(
            prompt=final_prompt,
            system_prompt=synthesis_prompt,
            role_context=synthesis_role,
            llm_profile=synthesis_llm
        )
    except Exception as e:
        final_result = f"Synthesis failed: {e}\n\nRaw team results:\n{synthesis_context}"
    
    elapsed = time.time() - t0
    
    return CrewRunResponse(
        success=True,
        result=final_result,
        meta={
            "orchestrator_id": orchestrator_id,
            "profile": profile,
            "team_name": team_name,
            "roles_count": len(team_members),
            "roles_success": len(successful_results),
            "roles_failed": len(failed_results),
            "elapsed_seconds": round(elapsed, 2),
            "parallel": parallel_roles
        }
    )


# === ENDPOINTS ===

@app.get("/health")
async def health():
    return {"status": "ok", "version": "2.0.0", "variant": "A-profiles"}


@app.get("/crew/agents")
async def list_agents():
    """List all agents with their profiles"""
    return get_all_agents_summary()


@app.get("/crew/teams")
async def list_teams():
    """List all teams (for backwards compatibility, returns profiles view)"""
    summary = get_all_agents_summary()
    result = {}
    
    for agent_id, info in summary.items():
        for profile_name in info.get("profiles", []):
            key = f"{agent_id}" if profile_name == "default" else f"{agent_id}.{profile_name}"
            settings = get_team_settings(agent_id, profile_name)
            deleg = get_delegation_config(agent_id, profile_name)
            members = get_team_members(agent_id, profile_name)
            
            result[key] = {
                "agent_id": agent_id,
                "profile": profile_name,
                "team_name": settings.get("team_name"),
                "member_count": len(members),
                "members": [m.get("role_context") for m in members],
                "parallel_roles": settings.get("parallel_roles"),
                "max_concurrency": settings.get("max_concurrency"),
                "has_delegation": deleg.get("enabled", False)
            }
    
    return result


@app.post("/crew/run", response_model=CrewRunResponse)
async def run_crew(request: CrewRunRequest):
    """Execute multi-role orchestration for an agent"""
    orchestrator_id = request.orchestrator_id
    
    if not is_orchestrator(orchestrator_id):
        raise HTTPException(status_code=404, detail=f"Agent {orchestrator_id} not found or not an orchestrator")
    
    return await run_crew_canonical(
        orchestrator_id=orchestrator_id,
        task=request.task,
        context=request.context,
        profile=request.profile
    )


@app.post("/crew/delegate")
async def delegate(request: DelegationRequest):
    """Delegate task to another top-level agent (for DAARWIZZ)"""
    result = await delegate_to_agent(
        orchestrator_id=request.orchestrator_id,
        target_agent_id=request.target_agent_id,
        task=request.task,
        context=request.context,
        hops_remaining=request.hops_remaining
    )
    
    if result is None:
        raise HTTPException(status_code=400, detail="Delegation failed or not allowed")
    
    return {"success": True, "result": result}


@app.on_event("startup")
async def startup():
    config = load_teams_config()
    version = config.get("version", "unknown")
    summary = get_all_agents_summary()
    
    total_profiles = sum(len(info.get("profiles", [])) for info in summary.values())
    
    logger.info(f"=== CrewAI Service v2.0.0 (Variant A) ===")
    logger.info(f"Config version: {version}")
    logger.info(f"Agents: {len(summary)}, Total profiles: {total_profiles}")
    
    for agent_id, info in summary.items():
        profiles = info.get("profiles", [])
        logger.info(f"  {agent_id}: {profiles}")