import os
import sys
import json
import shutil
import unittest
from datetime import datetime as _dt
from PIL import Image

# Ensure tour-comfy is in import path so we can import the FastAPI app and database module
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from fastapi.testclient import TestClient
import database
from edit_api import app, _load_output_dir

class TestAPIRegression(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # Determine output directory
        cls.output_dir = _load_output_dir()
        cls.client = TestClient(app)
        
        # Unique mock identifiers to avoid conflicts
        cls.test_ref_filename = "test_regression_ref_image_123.png"
        cls.test_other_filename = "test_regression_other_image_123.png"
        cls.group_id = "test_regression_group_123"
        cls.test_ref_path = os.path.join(cls.output_dir, cls.test_ref_filename)
        cls.test_other_path = os.path.join(cls.output_dir, cls.test_other_filename)
        
        # Ensure any leftover artifacts from past failed runs are cleaned
        cls._cleanup_database()
        cls._cleanup_files()

        # Create dummy test files (100x100 pixels, RGB format)
        img = Image.new("RGB", (100, 100), color=(255, 0, 0))
        img.save(cls.test_ref_path, "PNG")
        
        img_other = Image.new("RGB", (100, 100), color=(0, 255, 0))
        img_other.save(cls.test_other_path, "PNG")

        # Mock metadata
        cls.name = "Regression Test Character"
        cls.tags = ["VISIBLE", "LIKE", "21+"]
        cls.embedding = [0.1] * 1024
        cls.clip_description = "A dummy regression test image description"
        cls.prompt = "masterpiece, high quality"
        cls.pose = "standing"
        cls.group_name = "Regression Group"
        cls.hidden = False
        cls.has_background = True
        cls.has_clothing = False
        cls.is_source = True
        cls.pose_description = "The model is standing and looking directly at the camera."
        cls.pose_skeleton = '{"keypoints": [1, 2, 3]}'

        # Insert reference image into the database
        database.upsert_person(
            cls.test_ref_filename,
            filepath=cls.test_ref_path,
            name=cls.name,
            group_id=cls.group_id,
            tags=cls.tags,
            embedding=cls.embedding,
            clip_description=cls.clip_description,
            prompt=cls.prompt,
            pose=cls.pose,
            sort_order=0,
            group_name=cls.group_name,
            hidden=cls.hidden,
            has_background=cls.has_background,
            has_clothing=cls.has_clothing,
            is_source=cls.is_source,
            pose_description=cls.pose_description,
            pose_skeleton=cls.pose_skeleton
        )

        # Insert second image in same group for reordering test
        database.upsert_person(
            cls.test_other_filename,
            filepath=cls.test_other_path,
            name=cls.name,
            group_id=cls.group_id,
            tags=cls.tags,
            embedding=cls.embedding,
            clip_description=cls.clip_description,
            prompt=cls.prompt,
            pose=cls.pose,
            sort_order=1,
            group_name=cls.group_name,
            hidden=cls.hidden,
            has_background=cls.has_background,
            has_clothing=cls.has_clothing,
            is_source=cls.is_source,
            pose_description=cls.pose_description,
            pose_skeleton=cls.pose_skeleton
        )
        
        # Track dynamically created files and database rows for cleanup
        cls.created_files = [cls.test_ref_path, cls.test_other_path]
        cls.created_db_rows = [cls.test_ref_filename, cls.test_other_filename]

    @classmethod
    def tearDownClass(cls):
        # Cleanup
        cls._cleanup_database()
        cls._cleanup_files()

    @classmethod
    def _cleanup_database(cls):
        # Safely delete any inserted test rows from the person database table
        conn = database.get_db_connection()
        cur = conn.cursor()
        try:
            # Delete any filenames starting with test_regression or containing ts_crop/ts_pad
            cur.execute("""
                DELETE FROM person 
                WHERE filename LIKE 'test_regression_%%' 
                   OR filename LIKE '%%_crop_test_regression_%%' 
                   OR filename LIKE '%%_pad_test_regression_%%'
                   OR group_id = %s
            """, (cls.group_id,))
            conn.commit()
        except Exception as e:
            print(f"Error cleaning up database: {e}")
            conn.rollback()
        finally:
            cur.close()
            database._put_db_connection(conn)

    @classmethod
    def _cleanup_files(cls):
        # Clean up files matching test patterns
        for f in os.listdir(cls.output_dir):
            if "test_regression_" in f or "_crop_" in f or "_pad_" in f:
                p = os.path.join(cls.output_dir, f)
                try:
                    if os.path.exists(p):
                        os.remove(p)
                except Exception as e:
                    print(f"Error removing file {p}: {e}")

    def assertEmbeddingEqual(self, val1, val2):
        if isinstance(val1, str):
            val1 = [float(x) for x in val1.strip("[]").split(",")]
        if isinstance(val2, str):
            val2 = [float(x) for x in val2.strip("[]").split(",")]
        self.assertEqual(len(val1), len(val2))
        for x, y in zip(val1, val2):
            self.assertAlmostEqual(x, y, places=4)

    def test_01_duplicate_copies_all_pose_and_meta_details(self):
        # Send duplicate request to FastAPI
        response = self.client.post(f"/images/{self.test_ref_filename}/duplicate")
        self.assertEqual(response.status_code, 200)
        res_data = response.json()
        self.assertEqual(res_data["status"], "success")
        
        new_filename = res_data["new_filename"]
        self.created_db_rows.append(new_filename)
        new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
        self.created_files.append(new_path)

        # Assert duplicate file exists on disk
        self.assertTrue(os.path.exists(new_path), f"Duplicated file {new_path} not found on disk")

        # Assert all metadata has been accurately duplicated in DB
        person = database.get_person(new_filename)
        self.assertIsNotNone(person, "Duplicated database entry not found")
        
        # Column mappings as in database.py get_person:
        # SELECT name, group_id, tags, embedding, clip_description, filepath,
        #        prompt, pose, sort_order, group_name, hidden, has_background, source_refs,
        #        has_clothing, is_source, pose_description, pose_skeleton
        self.assertEqual(person[0], self.name)
        self.assertEqual(person[1], self.group_id)
        
        # Tags (assert LIKE and 21+ are preserved)
        tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
        self.assertIn("LIKE", tags_list)
        self.assertIn("21+", tags_list)
        self.assertIn("VISIBLE", tags_list)
        
        # Verify 21+ is just a standard tag, not a safety blocker
        self.assertIn("21+", tags_list)

        self.assertEmbeddingEqual(person[3], self.embedding)
        self.assertEqual(person[4], self.clip_description)
        self.assertEqual(person[5], new_path)
        self.assertEqual(person[6], self.prompt)
        self.assertEqual(person[7], self.pose)
        self.assertEqual(person[9], self.group_name)
        self.assertEqual(person[10], self.hidden)
        self.assertEqual(person[11], self.has_background)
        
        # source_refs should refer to original
        source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
        self.assertIn(self.test_ref_filename, source_refs)
        
        self.assertEqual(person[13], self.has_clothing)
        self.assertEqual(person[14], self.is_source)
        self.assertEqual(person[15], self.pose_description)
        self.assertEqual(person[16], self.pose_skeleton)

    def test_02_crop_as_copy_copies_all_pose_and_meta_details(self):
        # Crop region: (10, 10, 90, 90), as_copy=True
        req_payload = {
            "x1": 10,
            "y1": 10,
            "x2": 90,
            "y2": 90,
            "as_copy": True
        }
        response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload)
        self.assertEqual(response.status_code, 200)
        res_data = response.json()
        self.assertEqual(res_data["status"], "success")
        
        new_filename = res_data["new_filename"]
        self.created_db_rows.append(new_filename)
        new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
        self.created_files.append(new_path)

        # Verify physical file existence and crop dimensions (should be 80x80)
        self.assertTrue(os.path.exists(new_path), f"Cropped file {new_path} not found on disk")
        cropped_img = Image.open(new_path)
        self.assertEqual(cropped_img.size, (80, 80))

        # Verify database entry has complete metadata
        person = database.get_person(new_filename)
        self.assertIsNotNone(person, "Cropped copy database entry not found")
        
        self.assertEqual(person[0], self.name)
        self.assertEqual(person[1], self.group_id)
        
        tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
        self.assertIn("LIKE", tags_list)
        self.assertIn("21+", tags_list)
        self.assertIn("VISIBLE", tags_list)
        
        self.assertEmbeddingEqual(person[3], self.embedding)
        self.assertEqual(person[4], self.clip_description)
        self.assertEqual(person[5], new_path)
        self.assertEqual(person[6], self.prompt)
        self.assertEqual(person[7], self.pose)
        self.assertEqual(person[9], self.group_name)
        self.assertEqual(person[10], self.hidden)
        self.assertEqual(person[11], self.has_background)
        
        source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
        self.assertIn(self.test_ref_filename, source_refs)
        
        self.assertEqual(person[13], self.has_clothing)
        self.assertEqual(person[14], self.is_source)
        self.assertEqual(person[15], self.pose_description)
        self.assertEqual(person[16], self.pose_skeleton)

    def test_03_crop_in_place_keeps_all_meta_information(self):
        # Crop region: (20, 20, 80, 80), as_copy=False (in-place)
        req_payload = {
            "x1": 20,
            "y1": 20,
            "x2": 80,
            "y2": 80,
            "as_copy": False
        }
        
        # First verify original size is 100x100
        orig_img = Image.open(self.test_ref_path)
        self.assertEqual(orig_img.size, (100, 100))

        response = self.client.post(f"/images/{self.test_ref_filename}/crop", json=req_payload)
        self.assertEqual(response.status_code, 200)
        res_data = response.json()
        self.assertEqual(res_data["status"], "success")
        self.assertEqual(res_data["new_filename"], self.test_ref_filename)

        # Verify physical file size updated to 60x60
        cropped_img = Image.open(self.test_ref_path)
        self.assertEqual(cropped_img.size, (60, 60))

        # Verify database entry has complete metadata untouched
        person = database.get_person(self.test_ref_filename)
        self.assertIsNotNone(person, "Database entry not found after in-place crop")
        
        self.assertEqual(person[0], self.name)
        self.assertEqual(person[1], self.group_id)
        
        tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
        self.assertIn("LIKE", tags_list)
        self.assertIn("21+", tags_list)
        
        self.assertEmbeddingEqual(person[3], self.embedding)
        self.assertEqual(person[4], self.clip_description)
        self.assertEqual(person[6], self.prompt)
        self.assertEqual(person[7], self.pose)
        self.assertEqual(person[9], self.group_name)
        self.assertEqual(person[10], self.hidden)
        self.assertEqual(person[11], self.has_background)
        self.assertEqual(person[13], self.has_clothing)
        self.assertEqual(person[14], self.is_source)
        self.assertEqual(person[15], self.pose_description)
        self.assertEqual(person[16], self.pose_skeleton)

    def test_04_pad_as_copy_copies_all_pose_and_meta_details(self):
        # Pad canvas: expand top and bottom by 10 pixels, as_copy=True
        req_payload = {
            "top": 10,
            "right": 0,
            "bottom": 10,
            "left": 0,
            "as_copy": True,
            "fill": "transparent",
            "outpaint": False
        }
        # Original size at this point is 60x60 (due to test_03 crop)
        response = self.client.post(f"/images/{self.test_ref_filename}/pad", json=req_payload)
        self.assertEqual(response.status_code, 200)
        res_data = response.json()
        self.assertEqual(res_data["status"], "success")
        
        new_filename = res_data["new_filename"]
        self.created_db_rows.append(new_filename)
        new_path = os.path.join(self.output_dir, os.path.basename(new_filename))
        self.created_files.append(new_path)

        # Verify physical file existence and pad dimensions (should be 60x80)
        self.assertTrue(os.path.exists(new_path), f"Padded file {new_path} not found on disk")
        padded_img = Image.open(new_path)
        self.assertEqual(padded_img.size, (60, 80))

        # Verify database entry has complete metadata
        person = database.get_person(new_filename)
        self.assertIsNotNone(person, "Padded copy database entry not found")
        
        self.assertEqual(person[0], self.name)
        self.assertEqual(person[1], self.group_id)
        
        tags_list = json.loads(person[2]) if isinstance(person[2], str) else person[2]
        self.assertIn("LIKE", tags_list)
        self.assertIn("21+", tags_list)
        self.assertIn("VISIBLE", tags_list)
        
        self.assertEmbeddingEqual(person[3], self.embedding)
        self.assertEqual(person[4], self.clip_description)
        self.assertEqual(person[5], new_path)
        self.assertEqual(person[6], self.prompt)
        self.assertEqual(person[7], self.pose)
        self.assertEqual(person[9], self.group_name)
        self.assertEqual(person[10], self.hidden)
        self.assertEqual(person[11], self.has_background)
        
        source_refs = json.loads(person[12]) if isinstance(person[12], str) else person[12]
        self.assertIn(self.test_ref_filename, source_refs)
        
        self.assertEqual(person[13], self.has_clothing)
        self.assertEqual(person[14], self.is_source)
        self.assertEqual(person[15], self.pose_description)
        self.assertEqual(person[16], self.pose_skeleton)

    def test_05_reorder_group_updates_sort_orders(self):
        # Verify original order
        order_rows_before = database.get_group_order(self.group_id)
        filenames_before = [r[0] for r in order_rows_before]
        self.assertIn(self.test_ref_filename, filenames_before)
        self.assertIn(self.test_other_filename, filenames_before)

        # Reverse the order of files and submit
        new_order = [self.test_other_filename, self.test_ref_filename]
        req_payload = {
            "filenames": new_order
        }
        
        response = self.client.post(f"/groups/{self.group_id}/order", json=req_payload)
        self.assertEqual(response.status_code, 200)
        res_data = response.json()
        self.assertEqual(res_data["group_id"], self.group_id)
        self.assertEqual(res_data["filenames"], new_order)

        # Get order from database and verify state reflects the updates
        order_rows_after = database.get_group_order(self.group_id)
        filenames_after = [r[0] for r in order_rows_after]
        
        # Verify custom sort orders are explicitly 0, 1
        self.assertEqual(filenames_after[:2], new_order)
        
        person_other = database.get_person(self.test_other_filename)
        person_ref = database.get_person(self.test_ref_filename)
        
        # Column 8 in database.get_person is sort_order
        self.assertEqual(person_other[8], 0)
        self.assertEqual(person_ref[8], 1)

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