import os
import sys
import time
import unittest
from fastapi.testclient import TestClient

# Ensure tour-comfy is in the import path
_HERE = os.path.dirname(os.path.abspath(__file__))
if _HERE not in sys.path:
    sys.path.append(_HERE)

from edit_api import app, _load_output_dir

class TestQwenBackendRegression(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.client = TestClient(app)
        cls.output_dir = _load_output_dir()
        
        # Real-data filenames specified in the issue description
        cls.real_img_1 = "20260618_053519_1_20260618_053458_image.png"
        cls.real_img_2 = "20260618_052537_0_20260618_052526_image.png"
        cls.wireframe_img = "up_20260628_104641_image.png"

        # Validate that they exist on disk to run real end-to-end tests
        cls.run_real_tests = True
        img1_path = os.path.join(cls.output_dir, cls.real_img_1)
        img2_path = os.path.join(cls.output_dir, cls.real_img_2)
        
        if not os.path.exists(img1_path) or not os.path.exists(img2_path):
            cls.run_real_tests = False
            print(f"[test] Real images not found at {img1_path} or {img2_path}. Running tests with fallback/mock assertion modes.")

    def test_multi_ref_qwen_endpoint(self):
        """Test multi-reference generation endpoint with specified real-data images."""
        if not self.run_real_tests:
            self.skipTest("Skipping real backend test due to missing real image artifacts on this node.")

        payload = {
            "filenames": [self.real_img_1, self.real_img_2],
            "prompt": "standing girl looking at camera, high quality, highly detailed",
            "seed": 42,
            "max_area": 512 * 512,
            "pad_outpaint": False
        }
        
        response = self.client.post("/multi-ref", json=payload)
        self.assertEqual(response.status_code, 200)
        
        data = response.json()
        self.assertIn("job_id", data)
        job_id = data["job_id"]
        
        # Poll status until done (limit to 300 seconds to prevent hanging if backend is offline)
        print(f"[test] Polling multi-ref job {job_id} until completion...")
        start_time = time.time()
        completed = False
        
        while time.time() - start_time < 300:
            status_resp = self.client.get(f"/batch/{job_id}")
            if status_resp.status_code == 200:
                job_data = status_resp.json()
                status = job_data.get("status")
                print(f"[test] Job status: {status} (done: {job_data.get('done')}/{job_data.get('total')})")
                if status == "done":
                    completed = True
                    break
                elif status in ["error", "cancelled"]:
                    break
            time.sleep(3)
            
        self.assertTrue(completed, "Multi-ref generation job did not complete successfully in time.")

    def test_scenery_qwen_endpoint(self):
        """Test scenery generation endpoint with specified wireframe scene image and prompt."""
        if not self.run_real_tests:
            self.skipTest("Skipping real backend test due to missing real image artifacts on this node.")

        payload = {
            "model_filename": self.real_img_1,
            "scene_image": self.wireframe_img,
            "prompt": "replace person from image 1 naturally with the person from Image 2. Keep the exact position of Image 1, and the exact person of Image 2",
            "seed": 42
        }
        
        response = self.client.post("/generate-scenery", json=payload)
        self.assertEqual(response.status_code, 200)
        
        data = response.json()
        self.assertIn("job_id", data)
        job_id = data["job_id"]
        
        # Poll status until done
        print(f"[test] Polling scenery job {job_id} until completion...")
        start_time = time.time()
        completed = False
        
        while time.time() - start_time < 300:
            status_resp = self.client.get(f"/batch/{job_id}")
            if status_resp.status_code == 200:
                job_data = status_resp.json()
                status = job_data.get("status")
                print(f"[test] Job status: {status} (done: {job_data.get('done')}/{job_data.get('total')})")
                if status == "done":
                    completed = True
                    break
                elif status in ["error", "cancelled"]:
                    break
            time.sleep(3)
            
        self.assertTrue(completed, "Scenery generation job did not complete successfully in time.")

if __name__ == "__main__":
    unittest.main()