"""
Cube Agent with streaming support.

Runs in an agentic loop until task completion, streaming thoughts,
tool calls, and results back to the caller.
"""

import os
import sys
import json
import asyncio
from pathlib import Path
from typing import AsyncGenerator, Optional, List, Dict, Any
from datetime import datetime
import uuid

import anthropic
from anthropic import RateLimitError, APIConnectionError

from .models import EventType, StreamEvent, ToolCall, ToolResult
from .sessions import (
    Session,
    get_or_create_session,
    get_session,
    list_sessions,
    delete_session,
    SESSIONS_DIR
)

# Add cube-mcp to path
AGENT_DIR = Path(__file__).parent.parent
sys.path.insert(0, str(AGENT_DIR.parent / "src"))

# Import cube-mcp tools
try:
    from cube_mcp.server import (
        cube_login as _cube_login,
        cube_list as _cube_list,
        cube_status as _cube_status,
        acr_get_token as _acr_get_token,
        acr_login as _acr_login,
        helm_deploy as _helm_deploy
    )
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False
    print("Warning: cube_mcp module not found")

# Configuration
LLM_CONFIG_PATH = AGENT_DIR / "llm_config.json"
MODEL = "claude-sonnet-4-20250514"
MAX_TOKENS = 8192
MAX_TURNS = 50


# ============================================================================
# LLM Client
# ============================================================================

def load_llm_config() -> dict:
    """Load LLM configuration."""
    if LLM_CONFIG_PATH.exists():
        try:
            return json.loads(LLM_CONFIG_PATH.read_text())
        except Exception as e:
            print(f"Warning: Failed to load llm_config.json: {e}")
    return {}


def create_client():
    """Create Anthropic client based on config."""
    config = load_llm_config()
    provider = config.get("provider", "anthropic_bedrock")

    if provider == "anthropic_bedrock":
        bedrock_cfg = config.get("anthropic_bedrock", {})
        region = bedrock_cfg.get("aws_region") or os.environ.get("AWS_REGION", "us-west-2")

        access_key = bedrock_cfg.get("aws_access_key_id") or os.environ.get("AWS_ACCESS_KEY_ID")
        secret_key = bedrock_cfg.get("aws_secret_access_key") or os.environ.get("AWS_SECRET_ACCESS_KEY")

        if access_key and secret_key:
            return anthropic.AnthropicBedrock(
                aws_access_key=access_key,
                aws_secret_key=secret_key,
                aws_region=region,
            ), bedrock_cfg.get("model", MODEL)
        else:
            return anthropic.AnthropicBedrock(
                aws_region=region,
            ), bedrock_cfg.get("model", MODEL)
    else:
        anthropic_cfg = config.get("anthropic", {})
        api_key = anthropic_cfg.get("api_key") or os.environ.get("ANTHROPIC_API_KEY")

        if not api_key:
            raise ValueError("ANTHROPIC_API_KEY not set")

        return anthropic.Anthropic(api_key=api_key), anthropic_cfg.get("model", MODEL)


# ============================================================================
# System Prompt
# ============================================================================

SYSTEM_PROMPT = """You are the Cube Agent, an assistant for managing EdgescaleAI Cube edge devices and deployments.

## Capabilities

1. **Cube Cluster Access** (via Teleport)
   - cube_login: Connect to Cube clusters (staging-int, staging-pre-prod, ironmountain-v4)
   - cube_list: List available Cubes
   - cube_status: Check Cube node health

2. **Helm Deployments** (via Apollo Container Registry)
   - acr_login: Authenticate to ACR
   - acr_get_token: Get OAuth2 token (debugging)
   - helm_deploy: Full deployment pipeline (build, package, push, publish)

## Tool Details

### cube_login
Connect to a Cube via Teleport SSO (opens browser).

### cube_list
List all Cubes accessible via Teleport.

### cube_status
Get detailed node status (conditions, capacity, resources).

### acr_login
Authenticate Docker and Helm to Apollo Container Registry.
Requires APOLLO_CLIENT and APOLLO_SECRET environment variables.

### helm_deploy
Full pipeline: Docker build → Helm package → ACR push → Apollo publish.

Options:
- chart_path: Path to chart (auto-detects if not specified)
- version: Exact version (e.g., "1.0.0")
- bump: "patch", "minor", or "major" (default: patch)
- app_version: Docker image tag
- skip_docker: Skip Docker build
- dry_run: Preview without changes

## Workflow Examples

**Connect to Cube:**
1. Use cube_login(cluster="staging-int")
2. Verify with cube_status()

**Deploy Chart:**
1. First do dry_run: helm_deploy(dry_run=True)
2. If satisfied: helm_deploy()

## Notes

- Teleport login opens a browser for SSO
- ACR operations require Apollo credentials
- Always offer dry_run for deployments
- Be helpful with setup if credentials are missing
"""


# ============================================================================
# Tool Definitions
# ============================================================================

TOOLS = [
    {
        "name": "cube_login",
        "description": "Login to a Cube cluster via Teleport (opens browser for SSO).",
        "input_schema": {
            "type": "object",
            "properties": {
                "cluster": {
                    "type": "string",
                    "description": "Cube cluster name (e.g., staging-int, staging-pre-prod)"
                },
                "proxy": {
                    "type": "string",
                    "description": "Teleport proxy URL (default: edgescaleai.teleport.sh)"
                }
            },
            "required": ["cluster"]
        }
    },
    {
        "name": "cube_list",
        "description": "List all available Cube clusters.",
        "input_schema": {"type": "object", "properties": {}, "required": []}
    },
    {
        "name": "cube_status",
        "description": "Get status of the connected Cube node.",
        "input_schema": {"type": "object", "properties": {}, "required": []}
    },
    {
        "name": "acr_get_token",
        "description": "Get OAuth2 token from Apollo Container Registry.",
        "input_schema": {"type": "object", "properties": {}, "required": []}
    },
    {
        "name": "acr_login",
        "description": "Login to Apollo Container Registry (Docker + Helm).",
        "input_schema": {"type": "object", "properties": {}, "required": []}
    },
    {
        "name": "helm_deploy",
        "description": "Deploy Helm chart: build Docker, package chart, push to ACR, publish to Apollo.",
        "input_schema": {
            "type": "object",
            "properties": {
                "chart_path": {"type": "string", "description": "Path to chart (auto-detected if not specified)"},
                "version": {"type": "string", "description": "Exact version (e.g., '1.0.0')"},
                "bump": {"type": "string", "enum": ["patch", "minor", "major"], "description": "Version bump type"},
                "app_version": {"type": "string", "description": "Docker image tag version"},
                "skip_docker": {"type": "boolean", "description": "Skip Docker build"},
                "dry_run": {"type": "boolean", "description": "Preview without changes"}
            },
            "required": []
        }
    }
]


# ============================================================================
# Tool Execution
# ============================================================================

def execute_tool(name: str, input: dict) -> str:
    """Execute a tool and return the result."""
    if not MCP_AVAILABLE:
        return "Error: cube_mcp module not installed"

    try:
        if name == "cube_login":
            return _cube_login(
                cluster=input["cluster"],
                proxy=input.get("proxy", "edgescaleai.teleport.sh")
            )
        elif name == "cube_list":
            return _cube_list()
        elif name == "cube_status":
            return _cube_status()
        elif name == "acr_get_token":
            return _acr_get_token()
        elif name == "acr_login":
            return _acr_login()
        elif name == "helm_deploy":
            return _helm_deploy(
                chart_path=input.get("chart_path"),
                version=input.get("version"),
                bump=input.get("bump", "patch"),
                app_version=input.get("app_version"),
                skip_docker=input.get("skip_docker", False),
                dry_run=input.get("dry_run", False)
            )
        else:
            return f"Unknown tool: {name}"
    except Exception as e:
        return f"Error executing {name}: {e}"


# ============================================================================
# Streaming Agent Loop
# ============================================================================

async def run_agent_streaming(
    message: str,
    session: Session
) -> AsyncGenerator[StreamEvent, None]:
    """
    Run the agent in a streaming loop.

    Yields StreamEvents for:
    - thinking: Agent's reasoning
    - tool_call: Tool being called
    - tool_result: Result of tool call
    - text: Final response text
    - done: Agent finished
    - error: An error occurred
    """
    try:
        client, model = create_client()
    except Exception as e:
        yield StreamEvent(event=EventType.ERROR, data={"message": str(e)})
        return

    # Add user message to session
    session.add_user_message(message)

    turn = 0
    while turn < MAX_TURNS:
        turn += 1

        yield StreamEvent(
            event=EventType.THINKING,
            data={"turn": turn, "message": f"Processing turn {turn}..."}
        )

        # Call Claude with retry
        response = None
        for attempt in range(5):
            try:
                response = client.messages.create(
                    model=model,
                    max_tokens=MAX_TOKENS,
                    system=SYSTEM_PROMPT,
                    tools=TOOLS,
                    messages=session.get_messages_for_api()
                )
                break
            except (RateLimitError, APIConnectionError) as e:
                wait = 2 ** attempt
                yield StreamEvent(
                    event=EventType.THINKING,
                    data={"message": f"Rate limited, retrying in {wait}s..."}
                )
                await asyncio.sleep(wait)
            except Exception as e:
                yield StreamEvent(event=EventType.ERROR, data={"message": str(e)})
                return

        if not response:
            yield StreamEvent(event=EventType.ERROR, data={"message": "Failed to get response"})
            return

        # Process response
        if response.stop_reason == "end_turn":
            # Agent is done - extract text
            final_text = ""
            for block in response.content:
                if hasattr(block, "text"):
                    final_text += block.text

            session.add_assistant_message(response.content)

            yield StreamEvent(event=EventType.TEXT, data={"content": final_text})
            yield StreamEvent(
                event=EventType.DONE,
                data={"session_id": session.session_id, "turns": turn}
            )
            return

        elif response.stop_reason == "tool_use":
            # Process tool calls
            assistant_content = response.content
            session.add_assistant_message(assistant_content)

            tool_results = []

            for block in assistant_content:
                if block.type == "text" and block.text:
                    # Stream any thinking text
                    yield StreamEvent(
                        event=EventType.THINKING,
                        data={"content": block.text}
                    )

                elif block.type == "tool_use":
                    # Stream tool call
                    yield StreamEvent(
                        event=EventType.TOOL_CALL,
                        data={
                            "name": block.name,
                            "input": block.input,
                            "id": block.id
                        }
                    )

                    # Execute tool
                    result = execute_tool(block.name, block.input)

                    # Stream tool result
                    yield StreamEvent(
                        event=EventType.TOOL_RESULT,
                        data={
                            "tool_use_id": block.id,
                            "name": block.name,
                            "content": result[:2000]  # Truncate for streaming
                        }
                    )

                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result
                    })

            session.add_tool_results(tool_results)

        else:
            yield StreamEvent(
                event=EventType.ERROR,
                data={"message": f"Unexpected stop reason: {response.stop_reason}"}
            )
            return

    yield StreamEvent(
        event=EventType.ERROR,
        data={"message": f"Reached maximum turns ({MAX_TURNS})"}
    )


# ============================================================================
# Health Check Helpers
# ============================================================================

def check_tool_status() -> dict:
    """Check status of required tools."""
    import subprocess

    tools = {}
    checks = {
        "tsh": ["tsh", "version"],
        "kubectl": ["kubectl", "version", "--client"],
        "helm": ["helm", "version", "--short"],
        "docker": ["docker", "--version"],
        "apollo-cli": ["which", "apollo-cli"]
    }

    for name, cmd in checks.items():
        try:
            result = subprocess.run(cmd, capture_output=True, text=True, timeout=5)
            tools[name] = result.returncode == 0
        except:
            tools[name] = False

    return tools


def check_apollo_credentials() -> bool:
    """Check if Apollo credentials are set."""
    return bool(os.environ.get("APOLLO_CLIENT") and os.environ.get("APOLLO_SECRET"))


def get_model_name() -> str:
    """Get the configured model name."""
    config = load_llm_config()
    provider = config.get("provider", "anthropic_bedrock")
    if provider == "anthropic_bedrock":
        return config.get("anthropic_bedrock", {}).get("model", MODEL)
    return config.get("anthropic", {}).get("model", MODEL)
