"""
Test version of inference.py that uses rule-based actions (no LLM)
to verify the structured logging format matches requirements exactly.
"""
import os
import requests
from typing import List, Optional

ENV_URL = "http://localhost:7860"
BENCHMARK = "data-cleaning-env"
MODEL_NAME = "rule-based-test"
TASKS = ["easy", "medium", "hard"]


# ── Structured Logging Functions (MANDATORY FORMAT) ─────────────────────
def log_start(task: str, env: str, model: str) -> None:
    """Emit [START] line at episode begin."""
    print(f"[START] task={task} env={env} model={model}", flush=True)


def log_step(step: int, action: str, reward: float, done: bool, error: Optional[str]) -> None:
    """Emit [STEP] line after each env.step() with exact format."""
    error_val = error if error else "null"
    done_val = str(done).lower()  # Must be lowercase: true/false
    print(
        f"[STEP] step={step} action={action} reward={reward:.2f} done={done_val} error={error_val}",
        flush=True,
    )


def log_end(success: bool, steps: int, score: float, rewards: List[float]) -> None:
    """Emit [END] line after episode completes. Always called even on exception."""
    rewards_str = ",".join(f"{r:.2f}" for r in rewards)
    success_val = str(success).lower()  # Must be lowercase: true/false
    print(
        f"[END] success={success_val} steps={steps} score={score:.3f} rewards={rewards_str}",
        flush=True,
    )


def get_rule_based_action(obs: dict) -> dict:
    """Rule-based action selection (no LLM needed for testing)."""
    issues = obs["issues_detected"]
    legal = obs["legal_actions"]
    data = obs["current_data"]
    first_col = list(data.keys())[0] if data else "id"
    
    # Default
    action = {"action_type": "skip", "column": first_col, "value": None}
    
    if "missing:age" in issues and "fill_missing" in legal:
        action = {"action_type": "fill_missing", "column": "age", "value": 30}
    elif "missing:salary" in issues and "fill_missing" in legal:
        action = {"action_type": "fill_missing", "column": "salary", "value": 60000}
    elif "missing:score" in issues and "fill_missing" in legal:
        action = {"action_type": "fill_missing", "column": "score", "value": 80}
    elif "missing:name" in issues and "fill_missing" in legal:
        action = {"action_type": "fill_missing", "column": "name", "value": "unknown"}
    elif "wrong_type:score" in issues and "fix_type" in legal:
        action = {"action_type": "fix_type", "column": "score", "value": None}
    elif "formatting:email" in issues and "fix_formatting" in legal:
        action = {"action_type": "fix_formatting", "column": "email", "value": None}
    elif "formatting:name" in issues and "fix_formatting" in legal:
        action = {"action_type": "fix_formatting", "column": "name", "value": None}
    elif "formatting:dept" in issues and "fix_category" in legal:
        action = {"action_type": "fix_category", "column": "dept", "value": None}
    
    return action


def run_task(task_name: str) -> float:
    """Run one full episode for a task and return score in [0, 1] range."""
    rewards: List[float] = []
    steps_taken = 0
    score = 0.0
    success = False
    
    log_start(task=task_name, env=BENCHMARK, model=MODEL_NAME)
    
    try:
        # Reset environment
        obs = requests.post(
            f"{ENV_URL}/reset",
            json={"task_name": task_name}
        ).json()

        # Run episode until done
        while not obs["done"]:
            steps_taken += 1
            
            # Get action (rule-based, no LLM)
            error_msg = None
            try:
                action = get_rule_based_action(obs)
                action_str = f"{action['action_type']}('{action['column']}',{action['value']})"
            except Exception as e:
                error_msg = str(e)
                # Fallback action on error
                action = {
                    "action_type": "skip",
                    "column": list(obs["current_data"].keys())[0] if obs["current_data"] else "id",
                    "value": None
                }
                action_str = f"skip('{action['column']}',null)"

            # Take step in environment
            try:
                obs = requests.post(
                    f"{ENV_URL}/step",
                    json=action
                ).json()
                reward = obs.get("reward", 0.0)
                done = obs.get("done", False)
            except Exception as e:
                error_msg = str(e)
                reward = 0.0
                done = True
                obs = {"done": True, "reward": 0.0}
            
            rewards.append(reward)
            log_step(step=steps_taken, action=action_str, reward=reward, done=done, error=error_msg)
            
            if done:
                break

        # Get final score from environment grader (must be in [0, 1])
        try:
            state = requests.get(f"{ENV_URL}/state").json()
            score = state.get("score", 0.0)
            # Clamp score to [0, 1] range as required
            score = min(max(score, 0.0), 1.0)
            # Success if score meets reasonable threshold
            success = score >= 0.5
        except Exception as e:
            score = 0.0
            success = False

    except Exception as e:
        # Even on exception, we must emit [END]
        pass
    
    finally:
        # Always emit [END] line, even on exception
        log_end(success=success, steps=steps_taken, score=score, rewards=rewards)
    
    return score


def main():
    """Main inference script entry point."""
    # Health check - verify server is running
    try:
        health = requests.get(f"{ENV_URL}/health", timeout=5).json()
        if health.get("status") != "ok":
            print(f"[ERROR] Server not healthy at {ENV_URL}", flush=True)
            return
    except Exception as e:
        print(f"[ERROR] Server not running at {ENV_URL}: {e}", flush=True)
        return

    # Run inference on all tasks
    scores = {}
    for task in TASKS:
        scores[task] = run_task(task)

    # Print summary (not part of structured output, just for convenience)
    print(f"\n{'='*60}", flush=True)
    print("BASELINE RESULTS SUMMARY", flush=True)
    print(f"{'='*60}", flush=True)
    for task, score in scores.items():
        print(f"  {task:10s}: {score:.3f}", flush=True)
    avg = sum(scores.values()) / len(scores) if scores else 0.0
    print(f"  {'average':10s}: {avg:.3f}", flush=True)
    print(f"{'='*60}", flush=True)


if __name__ == "__main__":
    main()