import os
import textwrap
from typing import List
from openai import OpenAI
from dotenv import load_dotenv

from client import TrafficEnv
from models import TrafficAction

load_dotenv()

# Required Env Variables
API_BASE_URL = os.environ.get("API_BASE_URL")
MODEL_NAME = os.environ.get("MODEL_NAME")
HF_TOKEN = os.environ.get("HF_TOKEN")

REWARD_MIN_EXCLUSIVE = 0.01
REWARD_MAX_EXCLUSIVE = 0.99


def observation_to_dict(observation: object) -> dict:
    if isinstance(observation, dict):
        return observation
    if hasattr(observation, "model_dump"):
        return observation.model_dump()
    return dict(getattr(observation, "__dict__", {}))


def clamp_open_unit_interval(value: float) -> float:
    return max(REWARD_MIN_EXCLUSIVE, min(REWARD_MAX_EXCLUSIVE, value))

def build_user_prompt(step: int, last_action: str, last_reward: float, history: List[str], obs: dict) -> str:
    history_block = "\n".join(history[-4:]) if history else "None"
    return textwrap.dedent(
        f"""
        Step: {step}
        Last action taken: {last_action!r}
        Last reward: {last_reward:.2f}
        Previous steps history:
        {history_block}
        
        Current State:
        Task Difficulty: {obs.get('task_difficulty')}
        Signal Phase: {obs.get('signal_phase')}
        Time Since Last Change: {obs.get('time_since_last_change')}
        Queue Lengths: {obs.get('queue_lengths')}
        Waiting Times: {obs.get('waiting_times')}
        Emergency Presence: {obs.get('emergency_presence')}
        Pedestrian Requests: {obs.get('pedestrian_requests')}
        
        Decide your next action carefully to optimize the traffic. Your action must simply be a raw string of either 'KEEP_PHASE' or 'SWITCH_PHASE'.
        """
    ).strip()

def run_agent():
    # If variables are missing in local dev, fallback for dry-run
    api_key = HF_TOKEN or "dummy"
    llm_base_url = API_BASE_URL or "https://api.openai.com/v1"
    model = MODEL_NAME or "gpt-4o"
    env_url = os.environ.get("OPENENV_BASE_URL", "http://127.0.0.1:8000")
    llm_timeout_s = float(os.environ.get("LLM_TIMEOUT_S", "45"))
    episode_step_guard = int(os.environ.get("EPISODE_STEP_GUARD", "200"))
    
    client = OpenAI(api_key=api_key, base_url=llm_base_url, timeout=llm_timeout_s)
    env = TrafficEnv(base_url=env_url).sync()

    with env:
        for episode in range(3):
            # Start
            try:
                result = env.reset()
                obs_data = observation_to_dict(result.observation)
            except Exception as e:
                print(f"Failed to reset environment: {e}")
                break

            task = obs_data.get("task_difficulty", f"task_{episode}")
            env_name = "traffic_rl"

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

            step = 0
            done = False
            rewards = []
            raw_rewards = []
            history = []
            last_action = "None"
            last_reward = 0.0
            success = False
            score = 0.0

            while not done and step < episode_step_guard:
                prompt = build_user_prompt(step, last_action, last_reward, history, obs_data)

                try:
                    # LLM call
                    messages = [
                        {"role": "system", "content": "You are a highly efficient traffic signal controller. Output only 'KEEP_PHASE' or 'SWITCH_PHASE'."},
                        {"role": "user", "content": prompt},
                    ]
                    completion = client.chat.completions.create(
                        model=model,
                        messages=messages,
                        max_tokens=20,
                        timeout=llm_timeout_s,
                    )
                    action = completion.choices[0].message.content.strip().upper()

                    action = "SWITCH_PHASE" if "SWITCH" in action else "KEEP_PHASE"

                    error_val = "null"
                except Exception as e:
                    action = "KEEP_PHASE"
                    error_val = str(e).replace('\n', ' ')

                step += 1
                last_action = action

                try:
                    # Step the environment through a persistent OpenEnv session.
                    result = env.step(TrafficAction(action_type=action))
                    obs_data = observation_to_dict(result.observation)

                    raw_reward = float(result.reward if result.reward is not None else obs_data.get("reward", 0.0))
                    reward = clamp_open_unit_interval(raw_reward)
                    done_val = bool(result.done if result.done is not None else obs_data.get("done", False))
                except Exception as e:
                    raw_reward = 0.0
                    reward = clamp_open_unit_interval(raw_reward)
                    done_val = True
                    error_val = str(e).replace('\n', ' ')

                raw_rewards.append(raw_reward)
                rewards.append(reward)
                last_reward = reward
                done = done_val

                # Formatting as required by the spec exactly
                done_str = str(done_val).lower()
                print(f"[STEP] step={step} action={action} reward={reward:.2f} done={done_str} error={error_val}", flush=True)
                history.append(f"Phase={obs_data.get('signal_phase')} -> Act:{action} -> Rwd:{reward:.2f}")

            if not done and step >= episode_step_guard:
                done = True
                guard_reward = REWARD_MIN_EXCLUSIVE
                print(
                    f"[STEP] step={step} action=KEEP_PHASE reward={guard_reward:.2f} done=true error=Episode step guard reached ({episode_step_guard})",
                    flush=True,
                )

            # Compute Score inside [0.0, 1.0]. Easy tasks map higher positive reward. Hard tasks map lower.
            meta = obs_data.get("metadata", {})
            cumulative = meta.get("cumulative_reward", sum(raw_rewards))
            # Hard cap for grading
            score = max(REWARD_MIN_EXCLUSIVE, min(REWARD_MAX_EXCLUSIVE, (cumulative + 500) / 1000.0))
            success = score >= 0.3

            rewards_str = ",".join(f"{r:.2f}" for r in rewards)
            print(f"[END] success={str(success).lower()} steps={step} score={score:.3f} rewards={rewards_str}", flush=True)

if __name__ == "__main__":
    run_agent()