import os
import openai
import time

# Set API key
os.environ["OPENAI_API_KEY"] = "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"
openai.api_key = os.getenv("OPENAI_API_KEY")

MAX_PROMPTS_PER_SESSION = 5
THROTTLE_SECONDS = 30
last_usage_time = {}

def codette_terminal(prompt, model_name, generate_image, generate_video, session_id, batch_size, video_steps, fps):
    if session_id not in chat_memory:
        chat_memory[session_id] = []

    if prompt.lower() in ["exit", "quit"]:
        chat_memory[session_id] = []
        yield "🧠 Codette signing off... Session reset.", None, None
        return

    # --- Usage limits for fine-tuned model only ---
    if model_name == "Codette Fine-Tuned (v9)":
        count = sum(1 for line in chat_memory[session_id] if line.startswith("🖋️ You >"))
        if count >= MAX_PROMPTS_PER_SESSION:
            yield "[🛑 Usage Limit] You've reached the max prompt limit (5) for this session.", None, None
            return

        now = time.time()
        if now - last_usage_time.get(session_id, 0) < THROTTLE_SECONDS:
            wait = int(THROTTLE_SECONDS - (now - last_usage_time[session_id]))
            yield f"[⏳ Throttle] Wait {wait}s before trying again.", None, None
            return
        last_usage_time[session_id] = now

    response_so_far = ""

    if model_name == "Codette Fine-Tuned (v9)":
        try:
            response = openai.ChatCompletion.create(
                model="ft:gpt-4.1-2025-04-14:raiffs-bits:codette-final:BO907H7Z",
                messages=[{"role": "user", "content": prompt}],
                temperature=0.7,
                max_tokens=256
            )
            output = response.choices[0].message.content.strip()
        except Exception as e:
            yield f"[OpenAI fine-tuned model error]: {e}", None, None
            return
    else:
        if model_name not in text_model_cache:
            try:
                text_model_cache[model_name] = pipeline(
                    "text-generation",
                    model=AVAILABLE_MODELS[model_name],
                    device=0 if device == "cuda" else -1
                )
            except Exception as e:
                yield f"[Text model error]: {e}", None, None
                return

        generator = text_model_cache[model_name]
        try:
            output = generator(prompt, max_length=100, do_sample=True, num_return_sequences=1)[0]['generated_text'].strip()
        except Exception as e:
            yield f"[Text generation error]: {e}", None, None
            return

    for char in output:
        response_so_far += char
        temp_log = chat_memory[session_id][:]
        temp_log.append(f"🖋️ You > {prompt}")
        temp_log.append(f"🧠 Codette > {response_so_far}")
        yield "\n".join(temp_log[-10:]), None, None
        time.sleep(0.01)

    chat_memory[session_id].append(f"🖋️ You > {prompt}")
    chat_memory[session_id].append(f"🧠 Codette > {output}")

    imgs = None
    if generate_image and image_enabled:
        try:
            result = image_generator(prompt, num_images_per_prompt=batch_size)
            imgs = result.images
        except Exception as e:
            response_so_far += f"\n[Image error]: {e}"

    vid = None
    if generate_video and video_enabled:
        try:
            result = video_pipeline(prompt, num_inference_steps=video_steps)
            frames = result.frames
            temp_video_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
            imageio.mimsave(temp_video_path, frames, fps=fps)
            vid = temp_video_path
        except Exception as e:
            response_so_far += f"\n[Video error]: {e}"

    yield "\n".join(chat_memory[session_id][-10:]), imgs, vid