import os
import json
from dotenv import load_dotenv
from openai import OpenAI
from client import RevOpsEnvClient

load_dotenv()
api_key = os.getenv("OPENAI_API_KEY", "")
is_nvidia = api_key.startswith("nvapi-")
base_url = "https://integrate.api.nvidia.com/v1" if is_nvidia else None
model_name = "meta/llama-3.3-70b-instruct" if is_nvidia else "gpt-4o-mini"

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

env = RevOpsEnvClient("http://localhost:8000")

def act_on_task(task_id: str):
    print(f"--- Running Baseline for {task_id} ---")
    st = env.reset(task_id)
    
    SYSTEM_PROMPT = """You are a RevOps specialist AI. You process incoming leads based on rules. 
Available action formats (JSON only):
- {"action_type": "enrich_lead"}
- {"action_type": "check_crm"}
- {"action_type": "update_lead_score", "score": 85}
- {"action_type": "route_to_rep", "rep_id": "rep_amer_mm"}
- {"action_type": "merge_with_account", "account_id": "acc_major_fin"}
- {"action_type": "flag_reengagement", "opportunity_id": "opp_101"}
- {"action_type": "disqualify", "disqualification_reason": "student"}

Constraints:
1. Always enrich emails without full context.
2. Read the observation to see available reps, ICP criteria, and CRM data.
3. If a lead belongs to a closed-lost opportunity or existing account, merge and check first!
4. Only route to exactly one rep. 
5. Emit ONLY valid JSON corresponding to the chosen action type.
"""

    done = False
    step_count = 0
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT}
    ]
    
    while not done and step_count < 15:
        step_count += 1
        obs = st.get("observation", {})
        if not isinstance(obs, str):
            obs = json.dumps(obs)
        messages.append({"role": "user", "content": f"OBSERVATION:\n{obs}"})

        response = client.chat.completions.create(
            model=model_name,
            messages=messages,
            response_format={ "type": "json_object" }
        )
        
        action_payload_str = response.choices[0].message.content
        action_payload = json.loads(action_payload_str)
        print("ACTION:", action_payload)
        
        st = env.step(action_payload)
        done = st.get("done", True)
        
        messages.append({
            "role": "assistant",
            "content": f"Took action: {action_payload_str}. Result feedback: {st.get('metadata', {}).get('message', '')}"
        })

        
    final_score = env.grader()
    print(f"Task {task_id} Final Grader Score: {final_score}\n")
    return final_score

if __name__ == "__main__":
    tasks = ["task_easy", "task_medium", "task_hard"]
    scores = {}
    for t in tasks:
        scores[t] = act_on_task(t)
    print("All Baseline Scores:", scores)