const express = require('express');
const { chromium } = require('playwright');
const tf = require('@tensorflow/tfjs-node');
const axios = require('axios');
const sharp = require('sharp');

const app = express();
app.use(express.json());

const HCAPTCHA_FRAME = "iframe[src*='hcaptcha.com/captcha']";
const HCAPTCHA_CHECKBOX = "iframe[src*='hcaptcha.com/checkbox']";
const CHALLENGE_CONTAINER = ".challenge-container";

let model = null;

function log(...args) {
  console.log('[hCaptcha API]', new Date().toISOString(), ...args);
}

function delay(ms) {
  return new Promise(resolve => setTimeout(resolve, ms));
}

function randomDelay(min, max) {
  return delay(Math.floor(Math.random() * (max - min + 1)) + min);
}

async function loadModel() {
  if (!model) {
    try {
      log('Loading ML model...');
      model = await tf.loadLayersModel('https://tfhub.dev/google/tfjs-model/imagenet/mobilenet_v2_100_224/classification/3/default/1', {
        fromTFHub: true
      });
      log('ML model loaded successfully');
    } catch (error) {
      log('Failed to load model, using fallback:', error.message);
      model = null;
    }
  }
  return model;
}

async function preprocessImage(imageBuffer) {
  const processedImage = await sharp(imageBuffer)
    .resize(224, 224)
    .toBuffer();
  
  const tensor = tf.node.decodeImage(processedImage, 3)
    .toFloat()
    .div(255.0)
    .expandDims(0);
  
  return tensor;
}

async function analyzeImageWithML(imageUrl) {
  try {
    const response = await axios.get(imageUrl, { responseType: 'arraybuffer' });
    const imageBuffer = Buffer.from(response.data);
    
    const tensor = await preprocessImage(imageBuffer);
    
    if (model) {
      const predictions = await model.predict(tensor);
      const predArray = await predictions.data();
      const topPredictions = Array.from(predArray)
        .map((prob, idx) => ({ class: idx, probability: prob }))
        .sort((a, b) => b.probability - a.probability)
        .slice(0, 5);
      
      tensor.dispose();
      predictions.dispose();
      
      return topPredictions;
    }
    
    tensor.dispose();
    return null;
  } catch (error) {
    log('ML analysis error:', error.message);
    return null;
  }
}

async function analyzeImageFeatures(frame, imageElement) {
  return await frame.evaluate((img) => {
    const canvas = document.createElement('canvas');
    const ctx = canvas.getContext('2d');
    canvas.width = img.width;
    canvas.height = img.height;
    ctx.drawImage(img, 0, 0);
    
    const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
    let brightness = 0;
    let edges = 0;
    const colors = new Set();
    let redDominance = 0;
    let blueDominance = 0;
    let greenDominance = 0;
    
    for (let i = 0; i < imageData.data.length; i += 4) {
      const r = imageData.data[i];
      const g = imageData.data[i + 1];
      const b = imageData.data[i + 2];
      brightness += (r + g + b) / 3;
      
      if (r > g && r > b) redDominance++;
      if (b > r && b > g) blueDominance++;
      if (g > r && g > b) greenDominance++;
      
      const colorKey = `${Math.floor(r/50)*50},${Math.floor(g/50)*50},${Math.floor(b/50)*50}`;
      colors.add(colorKey);
      
      if (i + canvas.width * 4 < imageData.data.length) {
        if (Math.abs(r - imageData.data[i + canvas.width * 4]) > 30) edges++;
      }
    }
    
    const totalPixels = imageData.data.length / 4;
    
    return {
      brightness: brightness / totalPixels,
      edges: edges,
      colorDiversity: colors.size,
      redDominance: redDominance / totalPixels,
      blueDominance: blueDominance / totalPixels,
      greenDominance: greenDominance / totalPixels
    };
  }, imageElement);
}

async function classifyImage(imageUrl, challengeText, features) {
  const mlPredictions = await analyzeImageWithML(imageUrl);
  
  let score = 0;
  const challenge = challengeText.toLowerCase();
  
  if (mlPredictions && mlPredictions.length > 0) {
    const topClass = mlPredictions[0].class;
    
    if (challenge.includes('vehicle') || challenge.includes('car') || challenge.includes('bus') || challenge.includes('truck')) {
      if ((topClass >= 400 && topClass <= 900) || features.edges > 1000) {
        score += 3;
      }
    } else if (challenge.includes('animal') || challenge.includes('cat') || challenge.includes('dog') || challenge.includes('bird')) {
      if ((topClass >= 0 && topClass <= 400) || features.colorDiversity > 80) {
        score += 3;
      }
    } else if (challenge.includes('plant') || challenge.includes('tree') || challenge.includes('flower')) {
      if (features.greenDominance > 0.3 || (topClass >= 900 && topClass <= 1000)) {
        score += 3;
      }
    } else if (challenge.includes('sky') || challenge.includes('cloud')) {
      if (features.blueDominance > 0.4 || features.brightness > 150) {
        score += 3;
      }
    } else if (challenge.includes('water') || challenge.includes('ocean') || challenge.includes('lake')) {
      if (features.blueDominance > 0.3 || features.brightness > 100) {
        score += 3;
      }
    }
  }
  
  if (challenge.includes('vehicle') || challenge.includes('car') || challenge.includes('bus')) {
    if (features.edges > 1200) score += 2;
    if (features.colorDiversity > 70) score += 1;
  } else if (challenge.includes('animal') || challenge.includes('cat') || challenge.includes('dog')) {
    if (features.colorDiversity > 90) score += 2;
    if (features.edges > 800) score += 1;
  } else if (challenge.includes('plant') || challenge.includes('tree')) {
    if (features.greenDominance > 0.25) score += 2;
    if (features.colorDiversity > 60) score += 1;
  } else if (challenge.includes('sky') || challenge.includes('water')) {
    if (features.brightness > 140) score += 2;
    if (features.blueDominance > 0.3) score += 1;
  } else if (challenge.includes('building') || challenge.includes('house')) {
    if (features.edges > 1500) score += 2;
    if (features.colorDiversity < 70) score += 1;
  }
  
  return score;
}

async function findPuzzleMatch(frame) {
  log("Analyzing puzzle piece position with ML");
  
  const puzzleBox = await frame.$eval('.challenge-image', el => {
    const rect = el.getBoundingClientRect();
    return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
  });
  
  const pieceBox = await frame.$eval('.draggable-piece', el => {
    const rect = el.getBoundingClientRect();
    return { x: rect.x, y: rect.y, width: rect.width, height: rect.height };
  });
  
  const gapX = await frame.evaluate(() => {
    const canvas = document.querySelector('.challenge-image');
    const ctx = canvas.getContext('2d');
    const imageData = ctx.getImageData(0, 0, canvas.width, canvas.height);
    
    let darkestX = 0;
    let minBrightness = 999;
    let maxEdgeCount = 0;
    let edgeX = 0;
    
    for (let x = 0; x < canvas.width; x += 3) {
      let columnBrightness = 0;
      let edgeCount = 0;
      
      for (let y = 0; y < canvas.height; y++) {
        const i = (y * canvas.width + x) * 4;
        const brightness = (imageData.data[i] + imageData.data[i+1] + imageData.data[i+2]) / 3;
        columnBrightness += brightness;
        
        if (y > 0) {
          const prevI = ((y-1) * canvas.width + x) * 4;
          const diff = Math.abs(brightness - (imageData.data[prevI] + imageData.data[prevI+1] + imageData.data[prevI+2]) / 3);
          if (diff > 40) edgeCount++;
        }
      }
      
      columnBrightness /= canvas.height;
      
      if (columnBrightness < minBrightness) {
        minBrightness = columnBrightness;
        darkestX = x;
      }
      
      if (edgeCount > maxEdgeCount) {
        maxEdgeCount = edgeCount;
        edgeX = x;
      }
    }
    
    return (darkestX + edgeX) / 2;
  });
  
  return {
    targetX: puzzleBox.x + gapX,
    targetY: puzzleBox.y + puzzleBox.height / 2,
    startX: pieceBox.x + pieceBox.width / 2,
    startY: pieceBox.y + pieceBox.height / 2
  };
}

async function humanMouseMove(page, startX, startY, endX, endY) {
  await page.mouse.move(startX, startY);
  await randomDelay(50, 150);
  await page.mouse.down();
  await randomDelay(100, 200);
  
  const steps = 15 + Math.floor(Math.random() * 10);
  const controlX = (startX + endX) / 2 + (Math.random() - 0.5) * 50;
  const controlY = (startY + endY) / 2 + (Math.random() - 0.5) * 50;
  
  for (let i = 1; i <= steps; i++) {
    const t = i / steps;
    const x = Math.pow(1 - t, 2) * startX + 2 * (1 - t) * t * controlX + Math.pow(t, 2) * endX;
    const y = Math.pow(1 - t, 2) * startY + 2 * (1 - t) * t * controlY + Math.pow(t, 2) * endY;
    
    await page.mouse.move(x + (Math.random() - 0.5) * 2, y + (Math.random() - 0.5) * 2);
    await randomDelay(30, 80);
  }
  
  await randomDelay(150, 300);
  await page.mouse.up();
  await randomDelay(200, 400);
}

async function solveDragDropChallenge(page, challengeFrame) {
  log("Solving drag & drop challenge");
  
  const position = await findPuzzleMatch(challengeFrame);
  await humanMouseMove(page, position.startX, position.startY, position.targetX, position.targetY);
  
  await randomDelay(500, 1000);
  return true;
}

async function solveImageChallenge(challengeFrame, challengeText) {
  log("Solving image selection challenge with ML:", challengeText);
  
  const images = await challengeFrame.$$('.task-image');
  log(`Found ${images.length} images`);
  
  const imageData = [];
  
  for (let i = 0; i < images.length; i++) {
    const img = images[i];
    const features = await analyzeImageFeatures(challengeFrame, img);
    const imageUrl = await img.evaluate(el => el.src);
    const score = await classifyImage(imageUrl, challengeText, features);
    
    imageData.push({ index: i, score, features });
    await randomDelay(100, 200);
  }
  
  imageData.sort((a, b) => b.score - a.score);
  
  const threshold = imageData[0].score * 0.7;
  const selectedIndices = imageData
    .filter(data => data.score >= threshold)
    .map(data => data.index);
  
  if (selectedIndices.length === 0) {
    const randomCount = Math.min(3, Math.floor(Math.random() * 3) + 1);
    for (let i = 0; i < randomCount; i++) {
      selectedIndices.push(imageData[i].index);
    }
  }
  
  log(`Selecting ${selectedIndices.length} images based on ML:`, selectedIndices);
  
  for (const index of selectedIndices) {
    await randomDelay(300, 700);
    await images[index].click();
    await randomDelay(200, 400);
  }
  
  await randomDelay(500, 1000);
  
  const submitButton = await challengeFrame.$('.button-submit');
  if (submitButton) {
    await submitButton.click();
    await randomDelay(1500, 2500);
  }
  
  return true;
}

async function solveCaptcha(url, options = {}) {
  const {
    waitTimeout = 30000,
    maxRetries = 3
  } = options;
  
  await loadModel();
  
  const browser = await chromium.launch({
    headless: true,
    args: [
      '--no-sandbox',
      '--disable-setuid-sandbox',
      '--disable-dev-shm-usage',
      '--disable-accelerated-2d-canvas',
      '--disable-gpu'
    ]
  });
  
  const context = await browser.newContext({
    viewport: { width: 1280, height: 720 },
    userAgent: 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36'
  });
  
  const page = await context.newPage();
  
  try {
    log(`Navigating to: ${url}`);
    await page.goto(url, { waitUntil: 'networkidle', timeout: waitTimeout });
    
    await page.waitForSelector(HCAPTCHA_CHECKBOX, { state: "attached", timeout: waitTimeout });
    
    const checkboxElement = await page.$(HCAPTCHA_CHECKBOX);
    if (!checkboxElement) {
      throw new Error("Could not find hCaptcha checkbox iframe");
    }
    
    const checkboxFrame = await checkboxElement.contentFrame();
    if (!checkboxFrame) {
      throw new Error("Could not find hCaptcha checkbox iframe content");
    }
    
    const checkboxRect = await checkboxElement.boundingBox();
    if (!checkboxRect) {
      throw new Error("Could not get checkbox position");
    }
    
    log("Clicking checkbox with mouse");
    const checkboxX = checkboxRect.x + checkboxRect.width / 2;
    const checkboxY = checkboxRect.y + checkboxRect.height / 2;
    
    await page.mouse.move(checkboxX, checkboxY);
    await randomDelay(100, 300);
    await page.mouse.click(checkboxX, checkboxY);
    await randomDelay(2000, 3000);
    
    const challengeElement = await page.$(HCAPTCHA_FRAME);
    if (!challengeElement) {
      log("No challenge appeared - likely passed");
      
      const cookies = await context.cookies();
      await browser.close();
      
      return {
        success: true,
        message: "No challenge required",
        cookies: cookies
      };
    }
    
    const challengeFrame = await challengeElement.contentFrame();
    if (!challengeFrame) {
      throw new Error("Could not find hCaptcha challenge iframe content");
    }
    
    await challengeFrame.waitForSelector(CHALLENGE_CONTAINER, { timeout: waitTimeout });
    
    let attempts = 0;
    let solved = false;
    
    while (!solved && attempts < maxRetries) {
      attempts++;
      log(`Attempt ${attempts}/${maxRetries}`);
      
      const challengeText = await challengeFrame.$eval('.prompt-text', el => el.textContent).catch(() => '');
      log("Challenge prompt:", challengeText);
      
      const isDragDrop = await challengeFrame.$('.draggable-piece') !== null;
      
      try {
        if (isDragDrop) {
          await solveDragDropChallenge(page, challengeFrame);
        } else {
          await solveImageChallenge(challengeFrame, challengeText);
        }
        
        await randomDelay(2000, 3000);
        
        const stillHasChallenge = await page.$(HCAPTCHA_FRAME) !== null;
        if (!stillHasChallenge) {
          solved = true;
          log("Challenge solved successfully with ML!");
        }
      } catch (error) {
        log(`Attempt ${attempts} failed:`, error.message);
        await randomDelay(1000, 2000);
      }
    }
    
    if (!solved) {
      await browser.close();
      return {
        success: false,
        message: "Could not solve hCaptcha after maximum retries"
      };
    }
    
    const cookies = await context.cookies();
    const token = await page.evaluate(() => {
      const responseInput = document.querySelector('[name="h-captcha-response"]');
      return responseInput ? responseInput.value : null;
    });
    
    await browser.close();
    
    return {
      success: true,
      message: "hCaptcha solved successfully with ML",
      token: token,
      cookies: cookies
    };
    
  } catch (error) {
    await browser.close();
    throw error;
  }
}

app.post('/solve', async (req, res) => {
  try {
    const { url, options } = req.body;
    
    if (!url) {
      return res.status(400).json({
        success: false,
        error: 'URL is required'
      });
    }
    
    log(`Received solve request for: ${url}`);
    
    const result = await solveCaptcha(url, options);
    
    res.json(result);
    
  } catch (error) {
    log('Error:', error.message);
    res.status(500).json({
      success: false,
      error: error.message
    });
  }
});

app.get('/health', (req, res) => {
  res.json({ status: 'ok', model: model ? 'loaded' : 'fallback', timestamp: new Date().toISOString() });
});

const PORT = process.env.PORT || 7860;

app.listen(PORT, '0.0.0.0', () => {
  log(`hCaptcha Solver API with ML running on port ${PORT}`);
});