# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Data models for the FixOS Environment.

FixOS is a deterministic OS troubleshooting simulation where AI agents
diagnose and fix system issues through commands.
"""

import json
from typing import Any, Dict, List
from openenv.core.env_server.types import Action, Observation
from pydantic import BaseModel, Field, field_validator


class ProcessInfo(BaseModel):
    """Information about a running process."""
    pid: int = Field(..., description="Process ID")
    name: str = Field(..., description="Process name")
    cpu_percent: float = Field(..., description="CPU usage percentage")
    memory_mb: float = Field(..., description="Memory usage in MB")


class ServiceInfo(BaseModel):
    """Information about a system service."""
    name: str = Field(..., description="Service name")
    status: str = Field(..., description="Service status: running, stopped, failed")
    config_valid: bool = Field(..., description="Whether service config is valid")


class FileInfo(BaseModel):
    """Information about a file."""
    path: str = Field(..., description="File path")
    size_bytes: int = Field(..., description="File size in bytes")
    content_preview: str = Field(default="", description="First 500 chars of content")


class LogEntry(BaseModel):
    """A log message."""
    timestamp: int = Field(..., description="Timestamp in steps")
    level: str = Field(..., description="Log level: error, warning, info")
    message: str = Field(..., description="Log message")


class FixOSAction(Action):
    """Action for FixOS - a command to execute."""
    
    command: str = Field(..., description="Command to execute (ps, top, df, status, logs, cat, edit, restart, kill, rm)")
    args: Dict[str, Any] = Field(default_factory=dict, description="Command arguments")

    @field_validator("args", mode="before")
    @classmethod
    def parse_args_dict(cls, value: Any) -> Dict[str, Any]:
        # OpenEnv web UI can submit args as a JSON string; coerce it into a dict.
        if value is None:
            return {}
        if isinstance(value, dict):
            return value
        if isinstance(value, str):
            text = value.strip()
            if not text:
                return {}

            candidates = [text]
            if not text.startswith("{") and not text.endswith("}"):
                candidates.append("{" + text + "}")

            for candidate in candidates:
                try:
                    parsed = json.loads(candidate)
                    if isinstance(parsed, dict):
                        return parsed
                except json.JSONDecodeError:
                    continue

        raise TypeError("args must be a dictionary or JSON object string")


class FixOSObservation(Observation):
    """Observation from FixOS environment."""
    
    command_output: str = Field(default="", description="Output from the last command")
    processes: List[ProcessInfo] = Field(default_factory=list, description="Current processes")
    services: List[ServiceInfo] = Field(default_factory=list, description="Current services")
    filesystem: List[FileInfo] = Field(default_factory=list, description="Tracked files")
    resources: Dict[str, float] = Field(default_factory=dict, description="Resource aggregates (cpu%, memory%, disk%)")
    logs: List[LogEntry] = Field(default_factory=list, description="Recent log entries")
    history: List[str] = Field(default_factory=list, description="Command history")
    task_id: str = Field(default="", description="Current task ID")
    task_difficulty: str = Field(default="", description="Task difficulty: easy, medium, hard")
    task_score: float = Field(default=0.0, description="Current task progress score in [0.0, 1.0]")
    is_success_step: bool = Field(default=False, description="Whether task is solved")
    remaining_steps: int = Field(default=0, description="Steps remaining in episode")