#!/usr/bin/env python3
"""FixOS inference - FINAL PASS (LLM + FORMAT + SAFE)"""

import json
import os
from typing import Any, Dict

from openai import OpenAI


# =========================
# SAFE OBJECT → DICT
# =========================
def to_dict(obj: Any) -> Dict:
    if obj is None:
        return {}
    if isinstance(obj, dict):
        return obj
    if hasattr(obj, "model_dump"):
        return obj.model_dump()
    if hasattr(obj, "dict"):
        return obj.dict()
    return {}


# =========================
# LOAD ENV
# =========================
def _load_env():
    try:
        from server.my_env_environment import FixOSEnvironment
        return FixOSEnvironment()
    except Exception:
        from my_env_environment import FixOSEnvironment
        return FixOSEnvironment()


# =========================
# LLM CLIENT (MANDATORY FIX)
# =========================
def get_llm():
    api_key = os.environ.get("API_KEY")
    base_url = os.environ.get("API_BASE_URL")

    if not api_key or not base_url:
        raise RuntimeError("Missing API_KEY or API_BASE_URL")

    return OpenAI(
        api_key=api_key,
        base_url=base_url
    )


# =========================
# LLM ACTION
# =========================
def get_action(llm, observation: Dict) -> Dict:
    try:
        prompt = f"""
You are an OS troubleshooting agent.

Return ONLY JSON:
{{"command": "...", "args": {{}}}}

Observation:
{json.dumps(observation)}
"""

        resp = llm.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": prompt}],
            temperature=0,
            max_tokens=100,
        )

        content = resp.choices[0].message.content.strip()

        start = content.find("{")
        end = content.rfind("}")
        if start >= 0 and end >= 0:
            return json.loads(content[start:end+1])

    except Exception:
        pass

    return {"command": "status", "args": {}}


# =========================
# EPISODE
# =========================
def run_episode(env, llm, episode_id: int):
    try:
        reset_data = to_dict(env.reset())
        obs = reset_data.get("observation", reset_data)

        task = obs.get("task_id", f"task_{episode_id}")

        print(f"[START] task={task}", flush=True)

        step_count = 0

        for step in range(1, 51):
            step_count = step

            action = get_action(llm, obs)

            result = to_dict(env.step(action))
            obs = result.get("observation", result)

            reward = float(result.get("reward", obs.get("reward", 0)))
            done = bool(result.get("done", obs.get("done", False)))

            score = float(obs.get("task_score", 0))
            score = max(0.0001, min(0.9999, score))

            print(
                f"[STEP] step={step} reward={reward:.4f} score={score:.4f} done={done}",
                flush=True
            )

            if done:
                break

        final_score = float(obs.get("task_score", 0))
        final_score = max(0.0001, min(0.9999, final_score))

        print(
            f"[END] task={task} score={final_score:.4f} steps={step_count}",
            flush=True
        )

        return final_score

    except Exception:
        print(f"[STEP] step=0 reward=0.0000 score=0.0001 done=True", flush=True)
        print(f"[END] task=error score=0.0001 steps=0", flush=True)
        return 0.0001


# =========================
# MAIN
# =========================
def main():
    try:
        env = _load_env()
        llm = get_llm()

        scores = []
        for i in range(5):
            score = run_episode(env, llm, i)
            scores.append(score)

        avg = sum(scores) / len(scores)
        avg = max(0.0001, min(0.9999, avg))

        print(json.dumps({"summary": {"fixos": avg}}), flush=True)

    except Exception as e:
        print(f"[START] task=fail", flush=True)
        print(f"[STEP] step=0 reward=0.0000 score=0.0001 done=True", flush=True)
        print(f"[END] task=fail score=0.0001 steps=0", flush=True)
        print(json.dumps({"summary": {"fixos": 0.0001}}), flush=True)


if __name__ == "__main__":
    main()