Update app.py

app.py

@@ -231,4 +231,166 @@ def process_batch_gemini(api_key, items, model_name):
             model=model_name,
             contents=prompt_parts,
             config=types.GenerateContentConfig(
-                response_mime_type="application/json"
+                response_mime_type="application/json"
+            )
+        )
+        
+        text = response.text.replace("```json", "").replace("```", "").strip()
+        
+        try:
+            return json.loads(text)
+        except json.JSONDecodeError as e:
+            # Handle Truncated JSON (Output Token Limit Exceeded)
+            logger.warning(f"JSON Parse Error (likely truncated response): {e}. Attempting repair...")
+            
+            # Repair Strategy: Find the last closing brace '}', discard everything after, and close the array ']'
+            last_object_idx = text.rfind("}")
+            if last_object_idx != -1:
+                repaired_text = text[:last_object_idx+1] + "]"
+                try:
+                    repaired_data = json.loads(repaired_text)
+                    logger.info(f"Successfully repaired JSON. Recovered {len(repaired_data)}/{len(items)} items.")
+                    return repaired_data
+                except json.JSONDecodeError:
+                    logger.error("JSON repair failed.")
+            
+            return None # Fail gracefully if repair is impossible
+
+    except Exception as e:
+        logger.error(f"Gemini API Error with key ...{api_key[-4:]}: {e}")
+        return None 
+
+# --- Main Endpoint ---
+
+@app.post("/api/analyze")
+async def analyze_subtitles(
+    srt_file: UploadFile = File(...),
+    media_files: list[UploadFile] = File(...),
+    api_keys: str = Form(...),
+    batch_size: int = Form(20),
+    model_name: str = Form("gemini-2.0-flash"), # Updated default model hint
+    compression_quality: float = Form(0.7)
+):
+    temp_dir = tempfile.mkdtemp()
+    try:
+        # Convert float quality (0.1-1.0) to integer (10-100) for PIL
+        pil_quality = max(10, min(100, int(compression_quality * 100)))
+
+        # 1. Read SRT
+        srt_content = (await srt_file.read()).decode('utf-8', errors='ignore')
+        srt_data = parse_srt(srt_content)
+        srt_data.sort(key=lambda x: x['startTimeMs'])
+
+        # 2. Process Media
+        images = []
+        for file in media_files:
+            file_path = os.path.join(temp_dir, file.filename)
+            with open(file_path, "wb") as f:
+                shutil.copyfileobj(file.file, f)
+            
+            if file.filename.lower().endswith('.rar'):
+                try:
+                    with rarfile.RarFile(file_path) as rf:
+                        rf.extractall(temp_dir)
+                except rarfile.RarCannotExec:
+                     raise HTTPException(status_code=500, detail="Unrar executable not found in container.")
+            elif file.filename.lower().endswith('.zip'):
+                with zipfile.ZipFile(file_path, 'r') as zf:
+                    zf.extractall(temp_dir)
+        
+        for root, _, files in os.walk(temp_dir):
+            for filename in files:
+                if filename.lower().endswith(('.jpg', '.jpeg', '.png', '.webp', '.bmp')):
+                    full_path = os.path.join(root, filename)
+                    ms = parse_filename_to_ms(filename)
+                    if ms is not None:
+                        with open(full_path, "rb") as f:
+                            raw_bytes = f.read()
+                            compressed = compress_image(raw_bytes, quality=pil_quality)
+                            if compressed:
+                                images.append({
+                                    "filename": filename,
+                                    "timeMs": ms,
+                                    "data": compressed
+                                })
+
+        images.sort(key=lambda x: x['timeMs'])
+        
+        # 3. Pair
+        pairs = []
+        for i in range(len(images)):
+            img = images[i]
+            srt = srt_data[i] if i < len(srt_data) else None
+            
+            if srt:
+                # Create Thumbnail (lower quality for UI speed)
+                thumb_bytes = compress_image(img['data'], quality=50, max_width=300)
+                thumb_b64 = base64.b64encode(thumb_bytes).decode('utf-8')
+                
+                pairs.append({
+                    "index": i,
+                    "image_data": img['data'],
+                    "expected_text": srt['text'],
+                    "srt_id": srt['id'],
+                    "srt_time": srt['time'],
+                    "filename": img['filename'],
+                    "thumb": f"data:image/jpeg;base64,{thumb_b64}",
+                    "status": "pending"
+                })
+
+        if not pairs:
+            return {"status": "error", "message": "No valid image/subtitle pairs found."}
+
+        # 4. Process Gemini
+        keys = [k.strip() for k in api_keys.split('\n') if k.strip()]
+        if not keys:
+            raise HTTPException(status_code=400, detail="No API Keys provided")
+
+        results_map = {}
+        batches = [pairs[i:i + batch_size] for i in range(0, len(pairs), batch_size)]
+        
+        def worker(batch_idx, batch):
+            key = keys[batch_idx % len(keys)]
+            return process_batch_gemini(key, batch, model_name)
+
+        with ThreadPoolExecutor(max_workers=len(keys)) as executor:
+            futures = [executor.submit(worker, i, b) for i, b in enumerate(batches)]
+            for future in futures:
+                res = future.result()
+                if res:
+                    for item in res:
+                        results_map[item['index']] = item
+
+        # 5. Build Output
+        final_output = []
+        for p in pairs:
+            analysis = results_map.get(p['index'])
+            status = "pending"
+            reason = ""
+            detected = ""
+            if analysis:
+                status = "match" if analysis['match'] else "mismatch"
+                reason = analysis.get('reason', '')
+                detected = analysis.get('detected_text', '')
+
+            final_output.append({
+                "id": p['index'],
+                "filename": p['filename'],
+                "thumb": p['thumb'],
+                "expected": p['expected_text'],
+                "detected": detected,
+                "status": status,
+                "reason": reason,
+                "srt_id": p['srt_id'],
+                "srt_time": p['srt_time']
+            })
+
+        return {"status": "success", "results": final_output}
+
+    except Exception as e:
+        logger.error(f"Server Error: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+    finally:
+        shutil.rmtree(temp_dir)
+
+app.mount("/", StaticFiles(directory="static", html=True), name="static")