import os
import time
import gradio as gr
from byteplussdkarkruntime import Ark
import requests

# Get API key from Hugging Face secret "Key" and CLEAN it
API_KEY = os.environ.get("Key", "").strip()
print(f"✅ Key loaded, length: {len(API_KEY)}")

# Initialize client with proxy
client = Ark(
    base_url="https://1hit.no/proxy/proxy.php",
    api_key=API_KEY,
    timeout=30.0,
    max_retries=3,
)

# Fresh new prompts
DEFAULT_PROMPTS = {
    "Cinematic Nature": "Majestic drone shot soaring above misty mountains at golden hour, sunlight breaking through clouds, cinematic 4k quality, smooth motion --duration 5 --camerafixed false",
    "Urban Exploration": "Dynamic drone flight through narrow alleyways of a neon-lit Tokyo street at night, rain-slicked surfaces reflecting lights, cyberpunk aesthetic --duration 5 --camerafixed false",
    "Action Sequence": "High-speed drone chasing a sports car through a winding coastal road, dramatic cliffside views, action movie style --duration 5 --camerafixed false",
    "Abstract Art": "Surreal drone flight through floating geometric shapes in a dreamlike void, pastel colors, ethereal atmosphere --duration 5 --camerafixed false",
    "Wildlife Documentary": "Drone following a herd of elephants across the African savanna at sunset, National Geographic style, majestic wildlife cinematography --duration 5 --camerafixed false",
    "Sports Highlight": "Drone racing alongside professional skiers carving through fresh powder in the Alps, high-energy sports broadcast style --duration 5 --camerafixed false",
    "Beach Paradise": "Drone gliding over turquoise waters and white sand beaches, palm trees swaying, tropical paradise aerial view --duration 5 --camerafixed false",
    "Ancient Temple": "Drone circling around ancient Angkor Wat temple at dawn, mystical atmosphere, historical documentary style --duration 5 --camerafixed false"
}

def poll_via_json(task_id):
    """Check io.json for task status"""
    json_url = "https://1hit.no/proxy/io.json"
    try:
        response = requests.get(json_url)
        data = response.json()
        if task_id in data:
            task_data = data[task_id]
            status = task_data.get('status')
            if status == 'succeeded':
                return task_data.get('video_url')
            elif status == 'failed':
                return "FAILED"
            else:
                return "PROCESSING"
    except:
        pass
    return None

def generate_video(image_path, prompt_text, progress=gr.Progress()):
    """Generate video with proper polling"""
    
    if not API_KEY:
        yield "❌ API Key not configured. Please add 'Key' secret.", None
        return
    
    if image_path is None:
        yield "⚠️ Please upload an image first", None
        return
    
    try:
        progress(0, desc="Preparing image...")
        
        # FOR TESTING: Use working image URL
        image_url = "https://ark-doc.tos-ap-southeast-1.bytepluses.com/seepro_i2v%20.png"
        print(f"Using image URL: {image_url}")
        
        yield "✅ Image ready! Creating video request...", None
        
        progress(0.2, desc="Creating request...")
        
        # Create task
        try:
            create_result = client.content_generation.tasks.create(
                model="seedance-1-5-pro-251215",
                content=[
                    {"type": "text", "text": prompt_text},
                    {"type": "image_url", "image_url": {"url": image_url}}
                ]
            )
        except Exception as e:
            yield f"❌ API Error: {str(e)}", None
            return
        
        task_id = create_result.id
        print(f"Task created: {task_id}")
        yield f"✅ Task created: {task_id}", None
        
        progress(0.3, desc="Polling for results...")
        
        # Poll for results - TRY SDK FIRST, then fallback to JSON
        attempts = 0
        max_attempts = 120
        used_fallback = False
        
        while attempts < max_attempts:
            try:
                # Try SDK method first
                get_result = client.content_generation.tasks.get(task_id=task_id)
                status = get_result.status
                
                if status == "succeeded":
                    progress(1.0, desc="Complete!")
                    video_url = get_result.content.video_url if hasattr(get_result, 'content') else None
                    yield "✅ Video generated successfully!", video_url
                    return
                    
                elif status == "failed":
                    error_msg = get_result.error if hasattr(get_result, 'error') else "Unknown error"
                    yield f"❌ Failed: {error_msg}", None
                    return
                else:
                    progress(0.3 + (attempts/max_attempts)*0.7, desc=f"Status: {status}")
                    yield f"⏳ Status: {status}... (attempt {attempts + 1})", None
                    time.sleep(2)
                    attempts += 1
                    
            except Exception as e:
                # If SDK fails, use JSON fallback
                if not used_fallback:
                    print("SDK polling failed, using JSON fallback...")
                    used_fallback = True
                
                # Poll JSON every 5 seconds
                result = poll_via_json(task_id)
                if result == "FAILED":
                    yield "❌ Task failed (via JSON)", None
                    return
                elif result and result != "PROCESSING":
                    yield "✅ Video generated successfully! (via JSON)", result
                    return
                else:
                    progress(0.3 + (attempts/max_attempts)*0.7, desc="Status: processing (JSON)")
                    yield f"⏳ Status: processing... (JSON fallback, attempt {attempts + 1})", None
                    time.sleep(5)
                    attempts += 1
        
        yield "⏰ Timeout after 2 minutes", None
        
    except Exception as e:
        yield f"❌ Error: {str(e)}", None

def update_prompt(choice):
    return DEFAULT_PROMPTS.get(choice, "")

# MANUAL POLLING FUNCTIONS
def manual_poll(task_id):
    """Manually poll a specific task and update io.json"""
    if not task_id:
        return "❌ Please enter a task ID"
    
    try:
        # Call proxy to poll this specific task
        url = f"https://1hit.no/proxy/proxy.php?task_id={task_id}"
        headers = {"Authorization": f"Bearer {API_KEY}"}
        
        response = requests.get(url, headers=headers)
        
        # Also fetch updated io.json
        io_response = requests.get("https://1hit.no/proxy/io.json")
        io_data = io_response.json()
        
        if task_id in io_data:
            status = io_data[task_id].get('status')
            video_url = io_data[task_id].get('video_url')
            
            result = f"✅ Task {task_id}\n"
            result += f"Status: {status}\n"
            if video_url:
                result += f"Video URL: {video_url}\n"
            if 'response' in io_data[task_id]:
                result += f"Full response: {io_data[task_id]['response']}\n"
            return result
        else:
            return f"❌ Task {task_id} not found in io.json"
            
    except Exception as e:
        return f"❌ Error: {str(e)}"

def get_raw_json():
    """Fetch raw io.json"""
    try:
        r = requests.get("https://1hit.no/proxy/io.json")
        return r.json()
    except:
        return {"error": "Could not fetch io.json"}

# Create the interface WITHOUT theme in Blocks constructor
with gr.Blocks(title="BytePlus Video Generator") as demo:
    
    gr.Markdown("# 🎥 BytePlus Video Generator")
    
    with gr.Row():
        if API_KEY:
            gr.Markdown("✅ **API Key:** Configured")
        else:
            gr.Markdown("❌ **API Key:** Not configured - please add 'Key' secret")
    
    # Create tabs properly in Gradio 6.5.1
    with gr.Tabs():
        with gr.TabItem("🎬 Generate Video"):
            with gr.Row():
                with gr.Column():
                    image_input = gr.Image(
                        label="Upload Starting Image",
                        type="filepath",
                        height=300
                    )
                    
                    shot_type = gr.Dropdown(
                        choices=list(DEFAULT_PROMPTS.keys()),
                        label="🎬 Shot Type",
                        value="Cinematic Nature"
                    )
                    
                    prompt = gr.Textbox(
                        label="📝 Custom Prompt",
                        lines=3,
                        value=DEFAULT_PROMPTS["Cinematic Nature"]
                    )
                    
                    shot_type.change(fn=update_prompt, inputs=shot_type, outputs=prompt)
                    
                    generate_btn = gr.Button("🚀 Generate Video", variant="primary", size="lg")
                
                with gr.Column():
                    status = gr.Textbox(label="Status", lines=6)
                    video_output = gr.Video(label="Generated Video")
            
            # Quick shot buttons
            gr.Markdown("### Quick Shot Select")
            with gr.Row():
                gr.Button("🏔️ Nature").click(fn=lambda: "Cinematic Nature", outputs=shot_type)
                gr.Button("🌃 City").click(fn=lambda: "Urban Exploration", outputs=shot_type)
                gr.Button("🏎️ Action").click(fn=lambda: "Action Sequence", outputs=shot_type)
                gr.Button("🎨 Abstract").click(fn=lambda: "Abstract Art", outputs=shot_type)
                gr.Button("🦁 Wildlife").click(fn=lambda: "Wildlife Documentary", outputs=shot_type)
                gr.Button("⛷️ Sports").click(fn=lambda: "Sports Highlight", outputs=shot_type)
        
        with gr.TabItem("🔧 Manual Polling"):
            gr.Markdown("### 🔧 Manual Polling & Debug")
            
            with gr.Row():
                with gr.Column():
                    gr.Markdown("#### Poll Specific Task")
                    task_id_input = gr.Textbox(
                        label="Task ID", 
                        placeholder="Enter task ID (cgt-...)"
                    )
                    poll_btn = gr.Button("🔄 Poll Now", variant="primary")
                    poll_result = gr.Textbox(label="Poll Result", lines=10)
                    
                    poll_btn.click(
                        fn=manual_poll, 
                        inputs=task_id_input, 
                        outputs=poll_result
                    )
                
                with gr.Column():
                    gr.Markdown("#### Current io.json")
                    refresh_btn = gr.Button("🔄 Refresh io.json", variant="secondary")
                    raw_json = gr.JSON(label="io.json Contents")
                    
                    refresh_btn.click(
                        fn=get_raw_json, 
                        outputs=raw_json
                    )
            
            gr.Markdown("""
            ### 📝 Instructions
            1. Copy a task ID from above (like `cgt-20260217072358-hszt9`)
            2. Paste it in the Task ID field
            3. Click "Poll Now" to force an update
            4. Check io.json to see if status changed
            
            This manually triggers the proxy to poll BytePlus for updates!
            """)
    
    # Connect generate button
    generate_btn.click(
        fn=generate_video,
        inputs=[image_input, prompt],
        outputs=[status, video_output]
    )

if __name__ == "__main__":
    # Move theme to launch()
    demo.launch(
        server_name="0.0.0.0",
        theme=gr.themes.Soft()
    )