"""
inference.py - AI Executive Operations Manager
Runs all 3 tasks using an LLM agent and prints structured logs.

Environment variables:
  API_BASE_URL  - OpenAI-compatible API base URL
  MODEL_NAME    - Model to use (e.g., "gpt-4o-mini", "meta-llama/Llama-3.1-8B-Instruct")
  HF_TOKEN      - API key / HuggingFace token

STDOUT format (strict):
  [START] task=<task_id> env=exec-ops model=<model>
  [STEP]  step=<n> action=<type>('<id>') reward=<0.00> done=<true|false> error=<msg|null>
  [END]   success=<true|false> steps=<n> score=<0.00> rewards=<r1,r2,...>
"""

import os
import json
import sys
import dotenv

# Force UTF-8 output on Windows
if sys.stdout.encoding != "utf-8":
    sys.stdout.reconfigure(encoding="utf-8")

from openai import OpenAI
from env import ExecOpsEnv, Action
from env.grader import grade

# -------------------------------------------------------
# Configuration from environment
# -------------------------------------------------------
dotenv.load_dotenv()
API_BASE_URL = os.getenv("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME = os.getenv("MODEL_NAME", "gpt-4o-mini")
HF_TOKEN = os.getenv("HF_TOKEN")
LOCAL_IMAGE_NAME = os.getenv("LOCAL_IMAGE_NAME")

client = OpenAI(base_url=API_BASE_URL, api_key=HF_TOKEN)

SYSTEM_PROMPT = """You are an expert AI assistant helping a startup CEO (Alex Rivera at NovaTech AI) manage their inbox efficiently.

Your job: Given the current environment state, choose ONE action for ONE email.

Triage principles:
1. Handle the highest PRIORITY * URGENCY emails first
2. Reply to items requiring personal CEO attention (investor relations, critical incidents, legal deadlines)
3. Delegate items that can be handled by team members (HR, finance, operations)
4. Schedule items requiring a meeting
5. Ignore truly trivial items only when critical items remain

You MUST respond with ONLY valid JSON, no explanation:
{"type": "reply|schedule|delegate|ignore", "email_id": "<id>"}

IMPORTANT: email_id MUST be one of the ids from the "unhandled_emails" list (e.g. "e1", "e2").
Do NOT use goal_id values (e.g. "g1", "g2") as the email_id.

Choose wisely - you have limited steps. Ignoring P4-P5 items is heavily penalized."""


def _log(msg: str) -> None:
    """Print a structured log line to stdout."""
    print(msg, flush=True)


def _err(msg: str) -> None:
    """Print diagnostic info to stderr - never pollutes stdout."""
    print(msg, file=sys.stderr, flush=True)


def _clean_error(msg: str) -> str:
    """Collapse error message to a single line safe for the log format."""
    return msg.replace("\n", " ").replace("\r", "").strip()[:120]


def parse_action_from_response(text: str) -> dict:
    """Extract JSON action from LLM response, handling markdown code blocks."""
    text = text.strip()
    if "```json" in text:
        text = text.split("```json")[1].split("```")[0].strip()
    elif "```" in text:
        parts = text.split("```")
        for part in parts[1::2]:
            part = part.strip()
            if part.startswith("{"):
                text = part
                break
    start = text.find("{")
    end = text.rfind("}") + 1
    if start >= 0 and end > start:
        text = text[start:end]
    return json.loads(text)


def get_fallback_action(obs: dict) -> Action:
    """Fallback: pick highest priority * urgency unhandled email and reply or delegate."""
    inbox = obs.get("inbox", [])
    unhandled = [e for e in inbox if not e.get("handled", False)]
    if not unhandled:
        return None
    best = max(unhandled, key=lambda e: e.get("priority", 1) * e.get("urgency", 0.5))
    action_type = "delegate" if best.get("priority", 1) <= 2 else "reply"
    return Action(type=action_type, email_id=best["id"])


def run_task(task_id: str) -> float:
    """Run a single task with the LLM agent. Returns final grade in [0, 1]."""
    _log(f"[START] task={task_id} env=exec-ops model={MODEL_NAME}")

    step_count = 0
    rewards: list[float] = []
    pending_step_log = None  # deferred so we can force done=true on the last step
    final_score = 0.0
    env = None

    try:
        env = ExecOpsEnv(task_id)
        obs = env.reset()
        max_steps = obs.get("max_steps", 10)

        while not obs.get("done", False) and step_count < max_steps:
            unhandled = [e for e in obs.get("inbox", []) if not e.get("handled", False)]
            if not unhandled:
                break

            # Flush previous buffered step - it is not the last, so done=false
            if pending_step_log is not None:
                _log(pending_step_log.replace("_DONE_", "false"))
                pending_step_log = None

            prompt_state = {
                "time": obs.get("time"),
                "steps_remaining": obs.get("steps_remaining", max_steps - step_count),
                "unhandled_emails": [
                    {
                        "id": e["id"],
                        "sender": e["sender"].split("<")[0].strip(),
                        "subject": e["subject"],
                        "priority": e["priority"],
                        "urgency": round(e["urgency"], 2),
                    }
                    for e in sorted(unhandled, key=lambda x: -(x["priority"] * x["urgency"]))
                ],
                "pending_goals": [
                    {"goal_id": g["id"], "description": g["description"], "priority": g["priority"]}
                    for g in obs.get("pending_goals", [])
                ],
            }

            # --- Get action from LLM (with fallback) ---
            step_error = None
            action = None
            try:
                response = client.chat.completions.create(
                    model=MODEL_NAME,
                    messages=[
                        {"role": "system", "content": SYSTEM_PROMPT},
                        {"role": "user", "content": f"Current state:\n{json.dumps(prompt_state, indent=2)}\n\nChoose your action (JSON only):"}
                    ],
                    temperature=0.1,
                    max_tokens=80,
                )
                action_text = response.choices[0].message.content
                action_data = parse_action_from_response(action_text)
                action = Action(type=action_data["type"], email_id=action_data["email_id"])

                valid_ids = {e["id"] for e in unhandled}
                if action.email_id not in valid_ids:
                    raise ValueError(f"Invalid email_id: {action.email_id}")

            except Exception as e:
                step_error = _clean_error(str(e))
                _err(f"[fallback] LLM error ({task_id} step {step_count+1}): {type(e).__name__}: {e}")
                action = get_fallback_action(obs)
                if action is None:
                    break

            # --- Execute step; retry with fallback on failure ---
            candidates = [action]
            fb = get_fallback_action(obs)
            if fb is not None and (action is None or fb.email_id != action.email_id):
                candidates.append(fb)

            executed = False
            for attempt in candidates:
                try:
                    result = env.step(attempt)
                    reward = result.get("reward", 0.0)
                    done = result.get("done", False)
                    obs = result.get("observation", result)
                    if "done" not in obs:
                        obs["done"] = done
                    step_count += 1
                    rewards.append(reward)
                    action_str = f"{attempt.type}('{attempt.email_id}')"
                    error_field = step_error if step_error else "null"
                    pending_step_log = (
                        f"[STEP] step={step_count} action={action_str} "
                        f"reward={reward:.2f} done=_DONE_ error={error_field}"
                    )
                    executed = True
                    break
                except Exception as e:
                    step_error = _clean_error(str(e))
                    _err(f"[fallback] Step error ({task_id} step {step_count+1}): {type(e).__name__}: {e}")

            if not executed:
                break

    except Exception as e:
        _err(f"[error] Unexpected error in task '{task_id}': {e}")

    finally:
        # Flush the final buffered step - always done=true
        if pending_step_log is not None:
            _log(pending_step_log.replace("_DONE_", "true"))

        try:
            final_score = grade(env._state) if env is not None else 0.0
        except Exception:
            final_score = 0.0
        success = "true" if final_score >= 0.5 else "false"
        rewards_str = ",".join(f"{r:.2f}" for r in rewards) if rewards else "0.00"
        _log(f"[END] success={success} steps={step_count} score={final_score:.3f} rewards={rewards_str}")

    return final_score


def main():
    tasks = ["easy", "medium", "hard"]
    scores = {}

    for task_id in tasks:
        try:
            score = run_task(task_id)
            scores[task_id] = score
        except Exception as e:
            _err(f"[error] Task '{task_id}' failed: {e}")
            scores[task_id] = 0.0
            _log(f"[END] success=false steps=0 score=0.000 rewards=0.00")

    return sum(scores.values()) / len(scores) if scores else 0.0


if __name__ == "__main__":
    avg = main()
    sys.exit(0)