"""
Baseline inference script for the OrgOS OpenEnv environment.
Runs all three workflows (A / B / C) and reports scores.

Required environment variables:
    API_BASE_URL   — LLM API endpoint (OpenAI-compatible)
    MODEL_NAME     — model identifier (default: gpt-4o-mini)
    HF_TOKEN       — API key for the LLM endpoint
    ENV_URL        — environment server URL (default: http://localhost:8000)

STDOUT FORMAT (OpenEnv spec):
    [START] task=<workflow_name> env=orgos-openenv model=<model>
    [STEP]  step=<n> action=<json> reward=<0.00> done=<true|false> error=<msg|null>
    [END]   task=<workflow_name> score=<0.00> steps=<n>
"""

import json
import os
import re
import sys
import time
from typing import AsyncGenerator, Dict, List, Optional

import httpx
from openai import OpenAI
from dotenv import load_dotenv
load_dotenv()
# ------------------------------------------------------------------
# Config
# ------------------------------------------------------------------

API_BASE_URL = os.environ.get("API_BASE_URL", "https://api.openai.com/v1")
MODEL_NAME   = os.environ.get("MODEL_NAME",   "gpt-4o-mini")
HF_TOKEN     = os.environ.get("HF_TOKEN",     "")
ENV_URL      = os.environ.get("ENV_URL",       "http://localhost:8000")

if not HF_TOKEN:
    print("[WARNING] HF_TOKEN is not set — LLM calls may fail.", file=sys.stderr)

llm_client = OpenAI(api_key=HF_TOKEN or "sk-placeholder", base_url=API_BASE_URL)

# ------------------------------------------------------------------
# System prompt
# ------------------------------------------------------------------

SYSTEM_PROMPT = """\
You are OrgOS Agent — an enterprise workflow automation agent.
You operate across four SaaS applications: Jira, Zendesk, Salesforce, and Workday.

Each turn you receive a JSON observation with:
  - workflow_goal    : the task you must complete
  - pending_steps    : remaining steps in the workflow
  - app_states       : current state of each app
  - schema_hints     : field renames in effect this episode (e.g. {"jira.priority": "severity"})
  - active_rules     : current SLA / approval thresholds
  - message          : feedback from the last action
  - current_score    : your cumulative score (0.001–0.999)

Respond ONLY with a valid JSON object — no markdown, no explanation.

Action format:
  {"app": "<app>", "operation": "<op>", "args": {...}}

Available apps and key operations:
  jira:       get_issue, create_issue, update_status, set_priority, assign_owner,
              add_label, link_zendesk_ticket, close_issue, list_issues
  zendesk:    get_ticket, acknowledge_ticket, set_urgency, assign_agent,
              escalate_to_jira, resolve_ticket, add_note, list_tickets,
              create_agent_profile
  salesforce: get_account, list_accounts, update_deal_stage, flag_churn_risk,
              assign_account_owner, log_interaction, get_opportunity
  workday:    get_employee, list_employees, provision_access, log_sla_event,
              request_budget_approval, create_onboarding_task, complete_task

CRITICAL RULES:
1. Read schema_hints FIRST — if "jira.priority" → "severity", use "severity" not "priority" in args.
2. Complete ALL pending_steps in order.
3. Do not repeat a successful action.
4. When an operation creates a new resource (e.g. create_issue returns issue_id "JIRA-051"),
   use THAT returned ID for all subsequent operations on that resource — not a pre-existing ID.
5. If an operation fails, read the message carefully and adapt.
6. Use list_* operations to discover record IDs when needed — never assume an ID.
7. Stop when pending_steps is empty or done=true.

Example actions:
  # Workflow A (bug fix) — discover then chain ticket_id and returned issue_id:
  {"app": "zendesk", "operation": "list_tickets", "args": {"state": "new"}}
  {"app": "zendesk", "operation": "acknowledge_ticket", "args": {"ticket_number": "<ticket_number from list_tickets>"}}
  {"app": "jira", "operation": "create_issue", "args": {"title": "Bug fix for <customer>", "linked_zendesk": "<ticket_number>"}}
  # → create_issue returns {"issue_id": "JIRA-051", ...} — use JIRA-051 for assign_owner below
  {"app": "jira", "operation": "assign_owner", "args": {"issue_id": "<issue_id from create_issue>", "assignee": "<engineer>"}}

  # Workflow B (onboarding) — find the pending employee, then thread employee_id + territory:
  {"app": "workday", "operation": "list_employees", "args": {"status": "pending"}}
  # → returns exactly one record. Capture employee_id (e.g. "EMP-NEW-001") AND territory (e.g. "west").
  {"app": "workday", "operation": "create_onboarding_task", "args": {"employee_id": "<employee_id from list>"}}
  {"app": "workday", "operation": "provision_access", "args": {"employee_id": "<employee_id from list>", "app_name": "jira"}}
  # SF: assign that employee as owner of an account in THEIR territory:
  {"app": "salesforce", "operation": "list_accounts", "args": {"territory": "<territory from list_employees>"}}
  {"app": "salesforce", "operation": "assign_account_owner", "args": {"account_id": "<account_id from list>", "owner": "<employee_id from list>"}}
  # Jira: assign an open issue to the new hire's employee_id:
  {"app": "jira", "operation": "list_issues", "args": {"status": "open"}}
  {"app": "jira", "operation": "assign_owner", "args": {"issue_id": "<issue_id from list>", "assignee": "<employee_id from list_employees>"}}

  # Workflow C (churn) — discover the at-risk account, then scope all queries to it:
  {"app": "salesforce", "operation": "list_accounts", "args": {"health": "red"}}
  {"app": "salesforce", "operation": "flag_churn_risk", "args": {"account_id": "<account_id from list>"}}
  {"app": "zendesk", "operation": "get_ticket", "args": {"customer_id": "<account_id from list>"}}
  {"app": "jira", "operation": "list_issues", "args": {"customer_id": "<account_id from list>"}}
  {"app": "salesforce", "operation": "assign_account_owner", "args": {"account_id": "<account_id from list>", "owner": "<engineer>"}}
"""

WORKFLOW_NAMES = {
    "A": "workflow-a-bug-fix",
    "B": "workflow-b-onboarding",
    "C": "workflow-c-churn-alert",
}

# ------------------------------------------------------------------
# OpenEnv stdout logging helpers
# ------------------------------------------------------------------

def log_start(task: str, env_name: str, model: str) -> None:
    print(f"[START] task={task} env={env_name} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    error_val = error if error else "null"
    print(
        f"[STEP] step={step} action={action} reward={reward:.4f} "
        f"done={str(done).lower()} error={error_val}",
        flush=True,
    )


def log_end(task_name: str, score: float, steps: int) -> None:
    safe_score = max(0.001, min(0.999, float(score)))
    print(f"[END] task={task_name} score={safe_score:.4f} steps={steps}", flush=True)


# ------------------------------------------------------------------
# HTTP helpers
# ------------------------------------------------------------------

def api_post(path: str, payload: dict = None) -> dict:
    url  = ENV_URL.rstrip("/") + path
    resp = httpx.post(url, json=payload or {}, timeout=30)
    resp.raise_for_status()
    return resp.json()


def api_get(path: str) -> dict:
    url  = ENV_URL.rstrip("/") + path
    resp = httpx.get(url, timeout=10)
    resp.raise_for_status()
    return resp.json()


# ------------------------------------------------------------------
# Observation formatter
# ------------------------------------------------------------------

def obs_to_text(obs: dict) -> str:
    lines = [
        f"current_score:  {obs['current_score']}",
        f"step_count:     {obs['step_count']}",
        f"workflow_id:    {obs['workflow_id']}",
        "",
        "=== WORKFLOW GOAL ===",
        obs["workflow_goal"],
        "",
        "=== PENDING STEPS ===",
        "\n".join(f"  - {s}" for s in obs["pending_steps"]) or "  (all steps complete!)",
        "",
        "=== SCHEMA HINTS (use these field names) ===",
        json.dumps(obs["schema_hints"], indent=2) if obs["schema_hints"] else "  (no drift — use canonical names)",
        "",
        "=== ACTIVE RULES ===",
        json.dumps(obs["active_rules"], indent=2),
        "",
        "=== LAST MESSAGE ===",
        obs["message"],
        "",
        "=== APP STATES ===",
    ]
    for app_name, view in obs.get("app_states", {}).items():
        lines.append(f"  [{app_name.upper()}]")
        lines.append(f"  {view}")
        lines.append("")
    if obs.get("rule_violations"):
        lines.append("=== RULE VIOLATIONS (fix these!) ===")
        for v in obs["rule_violations"]:
            lines.append(f"  ⚠  {v}")
        lines.append("")
    return "\n".join(lines)


# ------------------------------------------------------------------
# Single-workflow inference loop
# ------------------------------------------------------------------

def run_workflow(workflow_id: str) -> float:
    task_name = WORKFLOW_NAMES.get(workflow_id, f"workflow-{workflow_id.lower()}")

    print(f"\n{'='*60}", file=sys.stderr)
    print(f"  Running Workflow {workflow_id}", file=sys.stderr)
    print(f"{'='*60}", file=sys.stderr)

    result  = api_post("/reset", {"workflow_id": workflow_id})
    obs     = result["observation"]
    history: List[dict] = []
    steps_taken = 0

    log_start(task=task_name, env_name="orgos-openenv", model=MODEL_NAME)

    try:
        for step_num in range(1, 60):
            if obs["done"]:
                break

            obs_text = obs_to_text(obs)
            history.append({"role": "user", "content": obs_text})

            # Trim history — keep last 20, ensure it starts with a user message
            if len(history) > 20:
                history = history[-20:]
            if history and history[0]["role"] != "user":
                history = history[1:]

            # Inject system prompt into first user message (Gemma/models without system role)
            messages_for_llm = list(history)
            if messages_for_llm:
                messages_for_llm[0] = {
                    "role": "user",
                    "content": SYSTEM_PROMPT + "\n\n---\n\n" + messages_for_llm[0]["content"],
                }

            try:
                response = llm_client.chat.completions.create(
                    model       = MODEL_NAME,
                    messages    = messages_for_llm,
                    temperature = 0.0,
                    max_tokens  = 300,
                )
                action_str = response.choices[0].message.content.strip()
            except Exception as exc:
                print(f"  Step {step_num}: LLM call failed: {exc}", file=sys.stderr)
                log_step(step_num, "null", 0.0, True, str(exc))
                break

            history.append({"role": "assistant", "content": action_str})

            # Parse action JSON
            action = None
            try:
                action = json.loads(action_str)
            except json.JSONDecodeError:
                m = re.search(r"\{.*\}", action_str, re.DOTALL)
                if m:
                    try:
                        action = json.loads(m.group())
                    except Exception:
                        pass

            if action is None:
                print(f"  Step {step_num}: Could not parse action JSON.", file=sys.stderr)
                log_step(step_num, action_str, -0.05, False, "json_parse_error")
                break

            action_label = json.dumps(action, separators=(",", ":"))
            print(
                f"  Step {step_num:2d} | score={obs['current_score']:.4f} | {action_label}",
                file=sys.stderr,
            )

            result      = api_post("/step", action)
            obs         = result["observation"]
            step_reward = result["reward"]
            done        = result["done"]
            error_msg   = (
                obs["message"]
                if obs.get("rule_violations") or step_reward < 0
                else None
            )

            print(f"           → {obs['message']}", file=sys.stderr)

            steps_taken = step_num
            log_step(
                step   = step_num,
                action = action_label,
                reward = step_reward,
                done   = done,
                error  = error_msg,
            )

            if done:
                break

            time.sleep(0.2)

    finally:
        final = obs.get("current_score", 0.001) if isinstance(obs, dict) else 0.001
        log_end(task_name=task_name, score=final, steps=steps_taken)

    final_score = obs["current_score"]
    wf_done     = not obs.get("pending_steps")
    print(
        f"\n  Workflow {workflow_id} final score: {final_score:.4f}  "
        f"steps: {obs['step_count']}  completed: {wf_done}",
        file=sys.stderr,
    )
    return final_score


# ------------------------------------------------------------------
# Async generator for SSE streaming from the UI
# ------------------------------------------------------------------

async def run_workflow_generator(
    workflow_id: str = "A",
    env_ref=None,
) -> AsyncGenerator[dict, None]:
    """
    Async generator that runs one inference episode and yields
    SSE-friendly event dicts for the dashboard UI.

    Each yielded dict has a "type" key:
      "reset"  — episode started
      "step"   — one action taken
      "done"   — episode ended
      "error"  — something went wrong
    """
    import asyncio

    if env_ref is None:
        # Fall back to HTTP if no direct env reference
        result = api_post("/reset", {"workflow_id": workflow_id})
    else:
        from models import OrgOSAction as _Action
        obs_obj = env_ref.reset(workflow_id=workflow_id)
        result  = {"observation": obs_obj.model_dump(), "reward": obs_obj.reward, "done": False}

    obs     = result["observation"]
    history: List[dict] = []

    yield {"type": "reset", "observation": obs, "workflow_id": workflow_id}
    await asyncio.sleep(0)

    for step_num in range(1, 60):
        if obs["done"]:
            break

        obs_text = obs_to_text(obs)
        history.append({"role": "user", "content": obs_text})
        if len(history) > 20:
            history = history[-20:]
        if history and history[0]["role"] != "user":
            history = history[1:]

        # Inject system prompt into first user message (Gemma/models without system role)
        messages_for_llm = list(history)
        if messages_for_llm:
            messages_for_llm[0] = {
                "role": "user",
                "content": SYSTEM_PROMPT + "\n\n---\n\n" + messages_for_llm[0]["content"],
            }

        try:
            response = llm_client.chat.completions.create(
                model       = MODEL_NAME,
                messages    = messages_for_llm,
                temperature = 0.0,
                max_tokens  = 300,
            )
            action_str = response.choices[0].message.content.strip()
        except Exception as exc:
            yield {"type": "error", "step": step_num, "message": str(exc)}
            break

        history.append({"role": "assistant", "content": action_str})

        action = None
        try:
            action = json.loads(action_str)
        except json.JSONDecodeError:
            m = re.search(r"\{.*\}", action_str, re.DOTALL)
            if m:
                try:
                    action = json.loads(m.group())
                except Exception:
                    pass

        if action is None:
            yield {"type": "error", "step": step_num, "message": "JSON parse error"}
            break

        if env_ref is None:
            result = api_post("/step", action)
        else:
            from models import OrgOSAction as _Action
            try:
                act     = _Action(**action)
                obs_obj = env_ref.step(act)
                result  = {
                    "observation": obs_obj.model_dump(),
                    "reward":      obs_obj.reward,
                    "done":        obs_obj.done,
                }
            except Exception as exc:
                yield {"type": "error", "step": step_num, "message": str(exc)}
                break

        obs         = result["observation"]
        step_reward = result["reward"]
        done        = result["done"]

        yield {
            "type":        "step",
            "step":        step_num,
            "action":      action,
            "observation": obs,
            "reward":      step_reward,
            "done":        done,
        }
        await asyncio.sleep(0)

        if done:
            break

    yield {
        "type":        "done",
        "final_score": obs.get("current_score", 0.001),
        "steps":       obs.get("step_count", step_num),
        "completed":   not obs.get("pending_steps"),
    }


# ------------------------------------------------------------------
# Main — run all three workflows sequentially
# ------------------------------------------------------------------

def main():
    print("OrgOS OpenEnv — Baseline Inference", file=sys.stderr)
    print(f"Model : {MODEL_NAME}", file=sys.stderr)
    print(f"Env   : {ENV_URL}", file=sys.stderr)

    try:
        health = api_get("/health")
        assert health.get("status") in ("ok", "healthy"), f"Unexpected status: {health}"
        print("Health check: OK\n", file=sys.stderr)
    except Exception as exc:
        print(f"[ERROR] Environment not reachable at {ENV_URL}: {exc}", file=sys.stderr)
        sys.exit(1)

    scores: Dict[str, float] = {}
    for wf_id in ["A", "B", "C"]:
        try:
            scores[f"workflow_{wf_id}"] = run_workflow(wf_id)
        except Exception as exc:
            print(f"[ERROR] Workflow {wf_id} failed: {exc}", file=sys.stderr)
            scores[f"workflow_{wf_id}"] = 0.001

    print("\n" + "="*60, file=sys.stderr)
    print("  BASELINE RESULTS", file=sys.stderr)
    print("="*60, file=sys.stderr)
    for k, v in scores.items():
        print(f"  {k}: {v:.4f}", file=sys.stderr)
    avg = round(sum(scores.values()) / len(scores), 4)
    print(f"  average: {avg:.4f}", file=sys.stderr)
    print("="*60, file=sys.stderr)

    with open("baseline_scores.json", "w") as f:
        json.dump({"scores": scores, "average": avg}, f, indent=2)
    print("\nScores written to baseline_scores.json", file=sys.stderr)


if __name__ == "__main__":
    main()