import os
import logging
from typing import TypedDict
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import StateGraph, END
# AsyncPostgresSaver imported lazily below (only when conn_string is provided)

from ops_agent.tools import (
    k8s_get_status, k8s_get_logs, k8s_restart_pod, k8s_scale_deployment,
    proxmox_get_status, proxmox_get_disk_space, proxmox_reboot_container, proxmox_clean_disk
)

logger = logging.getLogger(__name__)

class OpsState(TypedDict):
    target_type: str        # 'kubernetes' or 'proxmox'
    target_id: str          # e.g., 'default/web-pod' or VMID '101'
    status_info: str        # status output
    logs: str               # logs output
    diagnosis: str          # LLM diagnosis of issues
    proposed_action: str    # Proposed action: 'restart_pod', 'scale_deployment', 'reboot_container', 'clean_disk', 'none'
    proposed_action_details: str # E.g., pod name, deployment name, or container ID
    is_approved: bool       # Has user approved it?
    action_result: str      # Outcome of the action execution

def get_llm():
    """Initializes LLM using environment variables, matching the Gumbo setup."""
    api_key = os.environ.get("LITELLM_API_KEY") or os.environ.get("OPENAI_API_KEY", "mock-api-key")
    base_url = os.environ.get("LITELLM_BASE_URL") or os.environ.get("OPENAI_BASE_URL", "http://localhost:4000")
    model = os.environ.get("OPS_LLM_MODEL") or os.environ.get("GUMBO_LLM_MODEL", "ollama-qwen")
    
    logger.info(f"Initializing ChatOpenAI with model={model}, base_url={base_url}")
    return ChatOpenAI(
        model=model,
        api_key=api_key,
        base_url=base_url,
        temperature=0.1
    )

async def _fetch_status_and_logs(state: OpsState) -> dict:
    target_type = state.get("target_type", "").lower()
    target_id = state.get("target_id", "")
    
    status_info = ""
    logs = ""
    
    if target_type == "kubernetes":
        # target_id pattern: "namespace/pod-name" or just "pod-name" (assumes default namespace)
        parts = target_id.split("/")
        if len(parts) == 2:
            ns, pod = parts[0], parts[1]
        else:
            ns, pod = "default", parts[0]
            
        status_info = k8s_get_status(ns, pod)
        logs = k8s_get_logs(ns, pod)
        
    elif target_type == "proxmox":
        # target_id is LXC VMID (e.g. "101")
        status_info = proxmox_get_status(target_id)
        logs = proxmox_get_disk_space(target_id)
        
    else:
        status_info = f"Unknown target type: {target_type}"
        logs = "No logs available."
        
    return {"status_info": status_info, "logs": logs}

async def _diagnose_issue(state: OpsState) -> dict:
    llm = get_llm()
    system_msg = SystemMessage(
        content=(
            "You are OpsAgent, a senior DevOps engineer and Proxmox/Kubernetes site reliability assistant. "
            "Analyze the provided status info and logs, identify the root cause of any problems, and return "
            "a concise Markdown diagnosis."
        )
    )
    user_content = (
        f"Target Type: {state['target_type']}\n"
        f"Target ID: {state['target_id']}\n\n"
        f"=== STATUS INFO ===\n{state['status_info']}\n\n"
        f"=== LOGS / RESOURCE USAGE ===\n{state['logs']}\n"
    )
    
    try:
        resp = await llm.ainvoke([system_msg, HumanMessage(content=user_content)])
        diagnosis = resp.content if isinstance(resp.content, str) else str(resp.content)
    except Exception as e:
        logger.error(f"LLM call failed in diagnosis: {e}")
        diagnosis = f"LLM error: Failed to diagnose log. Raw error: {str(e)}"
        
    return {"diagnosis": diagnosis}

async def _propose_remediation(state: OpsState) -> dict:
    llm = get_llm()
    system_msg = SystemMessage(
        content=(
            "You are OpsAgent. Based on the diagnosis, propose a specific remediation action from the following list:\n"
            "- restart_pod (Kubernetes pods crash/hangs)\n"
            "- scale_deployment (K8s capacity issues)\n"
            "- reboot_container (Proxmox container stopped/unresponsive)\n"
            "- clean_disk (Proxmox/LXC container running out of disk space: Use% >= 90%)\n"
            "- none (No actions required, target is running fine)\n\n"
            "Respond ONLY with a JSON object: "
            '{"action": "restart_pod"|"scale_deployment"|"reboot_container"|"clean_disk"|"none", '
            '"details": "brief reason or extra target parameter"}'
        )
    )
    user_content = f"Diagnosis:\n{state['diagnosis']}"
    
    action = "none"
    details = "No action proposed."
    
    try:
        resp = await llm.ainvoke([system_msg, HumanMessage(content=user_content)])
        content = resp.content if isinstance(resp.content, str) else str(resp.content)
        # Parse JSON
        import json
        # Strip potential markdown fences
        clean_content = content.replace("```json", "").replace("```", "").strip()
        data = json.loads(clean_content)
        action = data.get("action", "none")
        details = data.get("details", "")
        
        # Dispatch WhatsApp Alert to phone if an anomaly/remediation is flagged
        if action != "none":
            try:
                from ops_agent.tools import send_whatsapp_alert
                alert_msg = (
                    f"⚠️ *OPSAGENT THREAT WARNING* ⚠️\n\n"
                    f"*Target:* {state['target_type'].upper()} ({state['target_id']})\n"
                    f"*Diagnosis:* {details}\n"
                    f"*Proposed Action:* {action.replace('_', ' ').upper()}\n\n"
                    f"Login to the Agentic OS Dashboard to authorize or reject this mitigation."
                )
                send_whatsapp_alert(alert_msg)
            except Exception as ex:
                logger.error(f"Could not dispatch WhatsApp alert: {ex}")
                
    except Exception as e:
        logger.error(f"Failed to propose action via LLM: {e}")
        
    return {"proposed_action": action, "proposed_action_details": details}

async def _execute_remediation(state: OpsState) -> dict:
    if not state.get("is_approved", False):
        return {"action_result": "Remediation action skipped: Approval was not provided."}
        
    action = state.get("proposed_action", "none")
    target_type = state.get("target_type", "")
    target_id = state.get("target_id", "")
    
    result = "No action executed."
    
    try:
        if target_type == "kubernetes":
            parts = target_id.split("/")
            ns = parts[0] if len(parts) == 2 else "default"
            pod = parts[1] if len(parts) == 2 else parts[0]
            
            if action == "restart_pod":
                result = k8s_restart_pod(ns, pod)
            elif action == "scale_deployment":
                # Default scale to 2 for demo purposes
                result = k8s_scale_deployment(ns, pod, 2)
            else:
                result = f"Action '{action}' is not supported for Kubernetes target."
                
        elif target_type == "proxmox":
            if action == "reboot_container":
                result = proxmox_reboot_container(target_id)
            elif action == "clean_disk":
                result = proxmox_clean_disk(target_id)
            else:
                result = f"Action '{action}' is not supported for Proxmox container VMID {target_id}."
        else:
            result = f"Unknown target type: {target_type}."
            
    except Exception as e:
        result = f"Error during execution: {str(e)}"
        
    return {"action_result": result}


def build_graph_builder() -> StateGraph:
    builder = StateGraph(OpsState)
    builder.add_node("fetch_status", _fetch_status_and_logs)
    builder.add_node("diagnose", _diagnose_issue)
    builder.add_node("propose_action", _propose_remediation)
    builder.add_node("execute_action", _execute_remediation)
    
    builder.set_entry_point("fetch_status")
    builder.add_edge("fetch_status", "diagnose")
    builder.add_edge("diagnose", "propose_action")
    builder.add_edge("propose_action", "execute_action")
    builder.add_edge("execute_action", END)
    
    return builder

async def run_ops_graph(target_type: str, target_id: str, is_approved: bool = False, conn_string: str = None) -> OpsState:
    builder = build_graph_builder()
    
    initial_state = {
        "target_type": target_type,
        "target_id": target_id,
        "status_info": "",
        "logs": "",
        "diagnosis": "",
        "proposed_action": "none",
        "proposed_action_details": "",
        "is_approved": is_approved,
        "action_result": ""
    }
    
    if conn_string:
        try:
            from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
        except ImportError:
            from langgraph_checkpoint_postgres.aio import AsyncPostgresSaver
        async with AsyncPostgresSaver.from_conn_string(conn_string) as checkpointer:
            await checkpointer.setup()
            graph = builder.compile(checkpointer=checkpointer)
            return await graph.ainvoke(
                initial_state,
                config={"configurable": {"thread_id": f"ops-{target_type}-{target_id}"}},
            )
    else:
        graph = builder.compile()
        return await graph.ainvoke(initial_state)
