import base64
import logging
import re
import requests
import cv2
import numpy as np
import pytesseract
from flask import Flask, render_template, jsonify
from threading import Lock

app = Flask(__name__)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

solved_captchas = []
lock = Lock()

CAPTCHA_URL = 'https://checkege.rustest.ru/api/captcha'
HEADERS = {
    'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
}

def order_points(pts):
    """
    Упорядочивает 4 точки в последовательности: верх-лево, верх-право, низ-право, низ-лево.
    """
    rect = np.zeros((4, 2), dtype="float32")
    s = pts.sum(axis=1)
    rect[0] = pts[np.argmin(s)]  # верх-лево
    rect[2] = pts[np.argmax(s)]  # низ-право
    diff = np.diff(pts, axis=1)
    rect[1] = pts[np.argmin(diff)]  # верх-право
    rect[3] = pts[np.argmax(diff)]  # низ-лево
    return rect

def find_text_region_contour(image):
    """
    Находит контур текстовой области более надежным способом.
    """
    # Создаем несколько версий для поиска контуров
    
    # Метод 1: Морфологические операции для объединения символов
    kernel_horizontal = cv2.getStructuringElement(cv2.MORPH_RECT, (25, 1))
    kernel_vertical = cv2.getStructuringElement(cv2.MORPH_RECT, (1, 10))
    
    # Расширяем по горизонтали чтобы соединить буквы
    dilated_h = cv2.dilate(image, kernel_horizontal, iterations=2)
    # Небольшое расширение по вертикали
    dilated = cv2.dilate(dilated_h, kernel_vertical, iterations=1)
    
    # Закрываем промежутки
    kernel_close = cv2.getStructuringElement(cv2.MORPH_RECT, (15, 8))
    closed = cv2.morphologyEx(dilated, cv2.MORPH_CLOSE, kernel_close, iterations=1)
    
    return closed

def correct_perspective_improved(image):
    """
    Улучшенная коррекция перспективы с более надежным поиском текстовой области.
    """
    logging.info("Запуск улучшенной коррекции перспективы...")
    
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()
    
    # Бинаризация если еще не сделана
    if len(np.unique(gray)) > 2:
        _, binary = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)
    else:
        binary = gray.copy()
    
    # Инвертируем если нужно (текст должен быть белым)
    if np.sum(binary == 255) < np.sum(binary == 0):
        binary = cv2.bitwise_not(binary)
    
    # Находим область текста
    text_region = find_text_region_contour(binary)
    
    # Находим контуры объединенной текстовой области
    contours, _ = cv2.findContours(text_region, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    if not contours:
        logging.warning("Контуры текстовой области не найдены.")
        return image
    
    # Берем самый большой контур
    largest_contour = max(contours, key=cv2.contourArea)
    
    # Проверяем минимальную площадь
    if cv2.contourArea(largest_contour) < 1000:
        logging.warning("Найденная текстовая область слишком мала.")
        return image
    
    # Аппроксимируем контур до прямоугольника
    epsilon = 0.02 * cv2.arcLength(largest_contour, True)
    approx = cv2.approxPolyDP(largest_contour, epsilon, True)
    
    # Если не получили 4 точки, используем минимальный ограничивающий прямоугольник
    if len(approx) != 4:
        rect = cv2.minAreaRect(largest_contour)
        box = cv2.boxPoints(rect)
        approx = np.int0(box)
    
    # Преобразуем к нужному формату
    if approx.shape[1] == 1:
        approx = approx.squeeze(1)
    
    # Упорядочиваем точки
    ordered_points = order_points(approx.astype("float32"))
    
    # Вычисляем размеры выходного изображения
    (tl, tr, br, bl) = ordered_points
    
    # Ширина: максимум из верхней и нижней стороны
    width_top = np.sqrt(((tr[0] - tl[0]) ** 2) + ((tr[1] - tl[1]) ** 2))
    width_bottom = np.sqrt(((br[0] - bl[0]) ** 2) + ((br[1] - bl[1]) ** 2))
    max_width = max(int(width_top), int(width_bottom))
    
    # Высота: максимум из левой и правой стороны
    height_left = np.sqrt(((bl[0] - tl[0]) ** 2) + ((bl[1] - tl[1]) ** 2))
    height_right = np.sqrt(((br[0] - tr[0]) ** 2) + ((br[1] - tr[1]) ** 2))
    max_height = max(int(height_left), int(height_right))
    
    # Проверяем адекватность размеров
    if max_width < 50 or max_height < 20:
        logging.warning(f"Неадекватные размеры области: {max_width}x{max_height}")
        return image
    
    # Добавляем небольшие отступы для лучшего распознавания
    padding = 10
    max_width += padding * 2
    max_height += padding * 2
    
    # Целевые точки (прямоугольник)
    dst_points = np.array([
        [padding, padding],                              # верх-лево
        [max_width - padding, padding],                  # верх-право
        [max_width - padding, max_height - padding],     # низ-право
        [padding, max_height - padding]                  # низ-лево
    ], dtype="float32")
    
    # Вычисляем матрицу трансформации и применяем
    transform_matrix = cv2.getPerspectiveTransform(ordered_points, dst_points)
    corrected = cv2.warpPerspective(
        image, 
        transform_matrix, 
        (max_width, max_height),
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(255, 255, 255)
    )
    
    logging.info(f"Коррекция перспективы применена. Размер: {max_width}x{max_height}")
    return corrected

def enhance_for_ocr(image):
    """
    Дополнительная обработка для улучшения OCR.
    """
    # Преобразуем в градации серого если нужно
    if len(image.shape) == 3:
        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    else:
        gray = image.copy()
    
    # Увеличиваем контраст
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    enhanced = clahe.apply(gray)
    
    # Бинаризация с адаптивным порогом
    binary = cv2.adaptiveThreshold(
        enhanced, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, 
        cv2.THRESH_BINARY, 11, 2
    )
    
    # Морфологическая очистка
    kernel = np.ones((2,2), np.uint8)
    cleaned = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, kernel, iterations=1)
    cleaned = cv2.morphologyEx(cleaned, cv2.MORPH_OPEN, kernel, iterations=1)
    
    return cleaned

def correct_perspective(image):
    """
    Основная функция коррекции перспективы.
    """
    # Применяем улучшенную коррекцию перспективы
    corrected = correct_perspective_improved(image)
    
    # Дополнительная обработка для OCR
    enhanced = enhance_for_ocr(corrected)
    
    return enhanced

def fetch_and_solve_captcha():
    try:
        logging.info("Получение новой капчи...")
        response = requests.get(CAPTCHA_URL, headers=HEADERS)
        response.raise_for_status()
        
        data = response.json()
        base64_image_data = data.get("Image")
        if not base64_image_data: return None

        image_bytes = base64.b64decode(base64_image_data)
        nparr = np.frombuffer(image_bytes, np.uint8)
        original_image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)

        # Увеличиваем масштаб для лучшего качества
        scale_factor = 3  # Увеличили с 2 до 3
        width = int(original_image.shape[1] * scale_factor)
        height = int(original_image.shape[0] * scale_factor)
        upscaled_image = cv2.resize(original_image, (width, height), interpolation=cv2.INTER_CUBIC)

        # Выделение синего текста
        hsv = cv2.cvtColor(upscaled_image, cv2.COLOR_BGR2HSV)
        lower_blue = np.array([90, 50, 50])
        upper_blue = np.array([130, 255, 255])
        mask = cv2.inRange(hsv, lower_blue, upper_blue)
        
        # Очистка маски
        kernel = np.ones((2, 2), np.uint8)
        cleaned_mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel, iterations=1)
        
        # ПРИМЕНЯЕМ УЛУЧШЕННУЮ КОРРЕКЦИЮ ПЕРСПЕКТИВЫ
        corrected_image = correct_perspective(cleaned_mask)
        
        # Настройки для Tesseract - более точные для цифр
        tesseract_config = r'--oem 3 --psm 8 -c tessedit_char_whitelist=0123456789'
        text = pytesseract.image_to_string(corrected_image, config=tesseract_config)
        recognized_text = re.sub(r'\s+', '', text).strip() or "Не распознано"
        logging.info(f"Распознано: {recognized_text}")

        # Кодируем изображения для отображения
        _, buffer_orig = cv2.imencode('.png', original_image)
        original_b64 = base64.b64encode(buffer_orig).decode('utf-8')
        
        _, buffer_proc = cv2.imencode('.png', corrected_image)
        processed_b64 = base64.b64encode(buffer_proc).decode('utf-8')

        return {
            "text": recognized_text,
            "original_b64": original_b64,
            "processed_b64": processed_b64
        }
    except Exception as e:
        logging.error(f"Произошла ошибка при решении капчи: {e}", exc_info=True)
        return None

@app.route('/')
def index():
    with lock:
        return render_template('index.html', captchas=list(solved_captchas))

@app.route('/solve', methods=['POST'])
def solve_new_captcha():
    new_captcha = fetch_and_solve_captcha()
    if new_captcha:
        with lock:
            solved_captchas.insert(0, new_captcha)
        return jsonify(new_captcha)
    return jsonify({"error": "Не удалось решить капчу"}), 500

if __name__ == '__main__':
    logging.info("Запуск приложения...")
    initial_captcha = fetch_and_solve_captcha()
    if initial_captcha:
        solved_captchas.append(initial_captcha)
    
    app.run(host='0.0.0.0', port=7860, debug=False)