from kubernetes import client, config
from openai import OpenAI
from pydantic import BaseModel, Field

# Initialize Kubernetes In-Cluster or Local Kubeconfig
try:
    config.load_incluster_config()
except config.ConfigException:
    config.load_kube_config()

core_api = client.CoreV1Api()
apps_api = client.AppsV1Api()
openai_client = OpenAI()

class CrashDiagnosis(BaseModel):
    root_cause: str = Field(description="Summary of the failure mechanism identified from logs.")
    recommended_action: str = Field(description="SCALE_MEMORY, RESTART, or ESCALATE_SRE")
    suggested_memory_mb: int = Field(description="New memory limit in MB if scaling is required.")
    confidence: float = Field(ge=0.0, le=1.0)

def triage_crashing_pod(namespace: str, pod_name: str, deployment_name: str) -> dict:
    # 1. Fetch previous crash logs
    try:
        logs = core_api.read_namespaced_pod_log(
            name=pod_name,
            namespace=namespace,
            tail_lines=150,
            previous=True,
        )
    except Exception as e:
        logs = f"Failed to fetch previous logs: {str(e)}"

    # 2. Query LLM for structured triage
    prompt = f"""
Analyze the Kubernetes container failure logs below:
Namespace: {namespace}
Pod: {pod_name}

Logs:
{logs}
"""
    response = openai_client.beta.chat.completions.parse(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": "You are a Kubernetes SRE expert. Diagnose pod crashes and recommend bounded actions."},
            {"role": "user", "content": prompt}
        ],
        response_format=CrashDiagnosis,
        temperature=0.1,
    )

    diagnosis = response.choices[0].message.parsed

    # 3. Apply bounded deployment memory patch if recommended
    if diagnosis.recommended_action == "SCALE_MEMORY" and diagnosis.suggested_memory_mb <= 4096:
        patch_body = {
            "spec": {
                "template": {
                    "spec": {
                        "containers": [
                            {
                                "name": deployment_name,
                                "resources": {
                                    "limits": {"memory": f"{diagnosis.suggested_memory_mb}Mi"}
                                }
                            }
                        ]
                    }
                }
            }
        }
        apps_api.patch_namespaced_deployment(
            name=deployment_name,
            namespace=namespace,
            body=patch_body
        )
        return {"status": "PATCHED", "diagnosis": diagnosis}

    return {"status": "ESCALATED", "diagnosis": diagnosis}