import requests
import config

class SmartRouter:
    def __init__(self):
        # Tracks what is currently loaded in the backend (Best Guess)
        self.lane_state = {
            "primary": None,    # URL: ...-generate-primary.modal.run
            "secondary": None   # URL: ...-generate-secondary.modal.run
        }

    def get_routing_plan(self, model_left_id, model_right_id):
        """
        Decides which model goes to which lane to minimize cold starts.
        Returns: (lane_for_left_model, lane_for_right_model)
        """
        primary_model = self.lane_state["primary"]
        secondary_model = self.lane_state["secondary"]

        # Score: 0 = Cache Hit (Good), 1 = Cache Miss (Bad)
        
        # Option A: Straight (Left -> Primary, Right -> Secondary)
        cost_straight = (0 if primary_model == model_left_id else 1) + \
                        (0 if secondary_model == model_right_id else 1)

        # Option B: Swapped (Left -> Secondary, Right -> Primary)
        cost_swapped = (0 if secondary_model == model_left_id else 1) + \
                       (0 if primary_model == model_right_id else 1)

        if cost_swapped < cost_straight:
            print(f"🔀 Smart Router: Swapping lanes to optimize cache!")
            # Update state for next time
            self.lane_state["secondary"] = model_left_id
            self.lane_state["primary"] = model_right_id
            return "secondary", "primary"
        else:
            print(f"⬇️ Smart Router: keeping straight lanes.")
            # Update state for next time
            self.lane_state["primary"] = model_left_id
            self.lane_state["secondary"] = model_right_id
            return "primary", "secondary"

# Create a global instance
router = SmartRouter()

# --- STEP 3: REWRITE call_modal_api FOR STREAMING ---
def call_modal_api(model_repo_id, prompt, lane):
    """
    Calls the Modal API on a specific lane and yields tokens as they arrive.
    This is now a GENERATOR.
    """
    if not model_repo_id:
        yield "Please select a model from the dropdown."
        return # Stop the generator
    
    if not config.MY_AUTH_TOKEN:
        yield "Error: `ARENA_AUTH_TOKEN` is not set on the Gradio server."
        return

    # Construct the URL based on the lane
    if lane == "primary":
        endpoint = f"{config.MODAL_BASE_URL}-generate-primary.modal.run"
    else:
        endpoint = f"{config.MODAL_BASE_URL}-generate-secondary.modal.run"

    print(f"🚀 Streaming from {model_repo_id} on [{lane.upper()}]...")

    headers = {
        "Content-Type": "application/json", 
        "Authorization": f"Bearer {config.MY_AUTH_TOKEN}"
    }
    payload = {"model_id": model_repo_id, "prompt": prompt}
    
    try:
        # stream=True is the magic.
        response = requests.post(
            endpoint, 
            json=payload, 
            timeout=300, 
            headers=headers, 
            stream=True
        )
        response.raise_for_status() 
        
        # Yield tokens as they arrive
        for chunk in response.iter_content(chunk_size=None, decode_unicode=True):
            if chunk:
                yield chunk
                
    except requests.exceptions.RequestException as e:
        if e.response and e.response.status_code == 401:
            yield "Error: Authentication failed. The token is invalid."
        elif e.response:
             # Try to get error detail from the streaming API
            try:
                error_detail = e.response.json().get("detail", str(e))
                yield f"API Error: {e.response.status_code} - {error_detail}"
            except:
                yield f"API Error: {e}"
        else:
            yield f"API Error: {e}"
    except Exception as e:
        yield f"An unexpected error occurred: {e}"