"""
香菇分析Web应用
基于Flask的Web界面，用于香菇直径和花纹分析
"""

from flask import Flask, render_template, request, jsonify
from werkzeug.utils import secure_filename
import os
import cv2
import numpy as np
import base64
from pathlib import Path
import logging
from typing import Optional, Tuple

try:
    from .integrated_analyzer import IntegratedMushroomAnalyzer
    from .paths import resolve_base_dir, ensure_app_dirs, resolve_dir
except ImportError:  # pragma: no cover - fallback for direct script usage
    from integrated_analyzer import IntegratedMushroomAnalyzer
    from paths import resolve_base_dir, ensure_app_dirs, resolve_dir


# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)

# 允许的文件扩展名
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp', 'webp'}


def allowed_file(filename: str) -> bool:
    """检查文件扩展名是否允许"""
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS


def image_to_base64(image):
    """将OpenCV图像转换为base64字符串"""
    if image is None:
        return None
    _, buffer = cv2.imencode('.jpg', image)
    return base64.b64encode(buffer).decode('utf-8')


def resolve_models_dir(models_dir: Optional[str], base_dir: Path) -> Path:
    if models_dir:
        return Path(models_dir).expanduser()

    env_dir = os.getenv("SMPAS_MODELS_DIR", os.getenv("MUSHROOM_MODELS_DIR"))
    if env_dir:
        return Path(env_dir).expanduser()

    repo_models = Path(__file__).resolve().parent.parent / "models"
    if repo_models.exists():
        return repo_models

    return base_dir / "models"


def resolve_model_paths(
    models_dir: Optional[str],
    yolo_model_path: Optional[str],
    sam_model_path: Optional[str],
    base_dir: Path
) -> Tuple[Path, Path]:
    if yolo_model_path:
        yolo_model = Path(yolo_model_path).expanduser()
    else:
        yolo_model = resolve_models_dir(models_dir, base_dir) / "yolo_mushroom.pt"

    if sam_model_path:
        sam_model = Path(sam_model_path).expanduser()
    else:
        sam_model = resolve_models_dir(models_dir, base_dir) / "sam_vit_b.pth"

    if not yolo_model.exists():
        raise FileNotFoundError(
            f"YOLO模型未找到: {yolo_model}\n"
            "请通过 --yolo-model 指定路径，或把模型放在 --models-dir 目录下。"
        )

    if not sam_model.exists():
        raise FileNotFoundError(
            f"SAM模型未找到: {sam_model}\n"
            "请通过 --sam-model 指定路径，或把模型放在 --models-dir 目录下。"
        )

    return yolo_model, sam_model


def create_app(
    base_dir: Optional[str] = None,
    models_dir: Optional[str] = None,
    data_dir: Optional[str] = None,
    yolo_model_path: Optional[str] = None,
    sam_model_path: Optional[str] = None,
    device: str = "cuda",
    enable_parallel: bool = True,
    max_workers: int = 4,
    apriltag_size_mm: float = 37.58,
    upload_dir: Optional[str] = None,
) -> Flask:
    template_dir = Path(__file__).parent / "templates"
    app = Flask(__name__, template_folder=str(template_dir))

    base_path = resolve_base_dir(base_dir)
    ensure_app_dirs(base_path)

    models_path = resolve_dir(base_path, models_dir, "SMPAS_MODELS_DIR", "models")
    data_path = resolve_dir(base_path, data_dir, "SMPAS_DATA_DIR", "data")
    uploads_path = resolve_dir(base_path, upload_dir, "SMPAS_UPLOAD_DIR", "uploads")
    models_path.mkdir(parents=True, exist_ok=True)
    data_path.mkdir(parents=True, exist_ok=True)
    uploads_path.mkdir(parents=True, exist_ok=True)

    app.config["MAX_CONTENT_LENGTH"] = 16 * 1024 * 1024  # 16MB max file size
    app.config["UPLOAD_FOLDER"] = str(uploads_path)
    app.config["BASE_DIR"] = str(base_path)
    app.config["MODELS_DIR"] = str(models_path)
    app.config["DATA_DIR"] = str(data_path)
    app.config["YOLO_MODEL"] = yolo_model_path
    app.config["SAM_MODEL"] = sam_model_path
    app.config["DEVICE"] = device
    app.config["ENABLE_PARALLEL"] = enable_parallel
    app.config["MAX_WORKERS"] = max_workers
    app.config["APRILTAG_SIZE_MM"] = apriltag_size_mm
    app.config["EXAMPLES_DIR"] = os.getenv(
        "SMPAS_EXAMPLES_DIR",
        os.getenv("MUSHROOM_EXAMPLES_DIR", str(data_path))
    )

    # 确保上传目录存在
    os.makedirs(app.config["UPLOAD_FOLDER"], exist_ok=True)

    analyzer = {"instance": None}

    def get_analyzer():
        """获取分析器实例（延迟加载）"""
        if analyzer["instance"] is None:
            logging.info("初始化香菇分析器...")
            yolo_model, sam_model = resolve_model_paths(
                app.config["MODELS_DIR"],
                app.config["YOLO_MODEL"],
                app.config["SAM_MODEL"],
                Path(app.config["BASE_DIR"]),
            )

            analyzer["instance"] = IntegratedMushroomAnalyzer(
                yolo_model_path=str(yolo_model),
                sam_model_path=str(sam_model),
                apriltag_size_mm=app.config["APRILTAG_SIZE_MM"],
                device=app.config["DEVICE"],
                enable_parallel=app.config["ENABLE_PARALLEL"],
                max_workers=app.config["MAX_WORKERS"],
            )
            logging.info("分析器初始化完成")

        return analyzer["instance"]

    @app.route('/')
    def index():
        """主页"""
        return render_template('index.html')

    @app.route('/analyze', methods=['POST'])
    def analyze():
        """分析上传的图像"""
        try:
            # 检查文件
            if 'file' not in request.files:
                return jsonify({'success': False, 'error': '未找到文件'}), 400

            file = request.files['file']
            if file.filename == '':
                return jsonify({'success': False, 'error': '未选择文件'}), 400

            if not allowed_file(file.filename):
                return jsonify({'success': False, 'error': '不支持的文件格式'}), 400

            # 保存文件
            filename = secure_filename(file.filename)
            filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(filepath)

            logging.info(f"处理上传文件: {filename}")

            # 获取分析器并处理
            analyzer_instance = get_analyzer()
            result = analyzer_instance.analyze_image(filepath, visualize=True)

            if result is None or not result.get('success'):
                error_msg = result.get('error', '分析失败') if result else '分析失败'
                return jsonify({'success': False, 'error': error_msg}), 500

            # 读取原始图像和可视化图像
            orig_image = cv2.imread(filepath)
            vis_image = result.get('visualization')

            # 准备响应数据
            response_data = {
                'success': True,
                'filename': filename,
                'apriltag_detected': bool(result.get('apriltag_detected', False)),
                'px_per_mm': float(round(result.get('px_per_mm', 0), 4)),
                'mushrooms': []
            }

            # 处理每个香菇的结果
            mushroom_results = result.get('mushrooms', [])
            for i, mushroom in enumerate(mushroom_results):
                mushroom_data = {
                    'id': int(i + 1),
                    'diameter_mm': float(round(mushroom.get('diameter_mm', 0), 4)),
                    'crack_ratio_percent': float(round(mushroom.get('crack_ratio_percent', 0), 2)),
                    'cap_area_mm2': float(round(mushroom.get('cap_area_mm2', 0), 2)),
                    'confidence': float(round(mushroom.get('detection_confidence', 0), 3)),
                    'pattern_type': str(mushroom.get('pattern_type', '未知')),
                    'pattern_category': str(mushroom.get('pattern_category', '未知')),
                    'white_pattern_ratio': float(round(mushroom.get('white_pattern_ratio', 0) * 100, 2)),
                    'brown_pattern_ratio': float(round(mushroom.get('brown_pattern_ratio', 0) * 100, 2)),
                    'flatness_score': float(round(mushroom.get('flatness_score', 0), 4)),
                    'suspected_ban_gu': bool(mushroom.get('suspected_ban_gu', False))
                }
                response_data['mushrooms'].append(mushroom_data)

            # 编码图像
            response_data['orig_image'] = image_to_base64(orig_image)
            response_data['vis_image'] = image_to_base64(vis_image)

            # 清理临时文件
            try:
                os.remove(filepath)
            except Exception:
                pass

            return jsonify(response_data)

        except Exception as e:
            logging.error(f"分析过程出错: {str(e)}", exc_info=True)
            return jsonify({'success': False, 'error': f'服务器错误: {str(e)}'}), 500

    @app.route('/examples', methods=['GET'])
    def get_examples():
        """获取示例图像列表"""
        examples_dir = app.config.get("EXAMPLES_DIR")
        if not examples_dir:
            return jsonify([])

        examples_path = Path(examples_dir)
        if not examples_path.exists():
            return jsonify([])

        examples = []
        for img_file in examples_path.glob('*.jpg'):
            examples.append({
                'name': img_file.name,
                'path': str(img_file)
            })

        # 限制返回数量
        return jsonify(examples[:10])

    @app.route('/analyze_example', methods=['POST'])
    def analyze_example():
        """分析示例图像"""
        try:
            data = request.get_json()
            image_path = data.get('path') if data else None

            if not image_path or not os.path.exists(image_path):
                return jsonify({'success': False, 'error': '示例图像不存在'}), 400

            logging.info(f"处理示例图像: {image_path}")

            # 获取分析器并处理
            analyzer_instance = get_analyzer()
            result = analyzer_instance.analyze_image(image_path, visualize=True)

            if result is None or not result.get('success'):
                error_msg = result.get('error', '分析失败') if result else '分析失败'
                return jsonify({'success': False, 'error': error_msg}), 500

            # 读取原始图像和可视化图像
            orig_image = cv2.imread(image_path)
            vis_image = result.get('visualization')

            # 准备响应数据
            response_data = {
                'success': True,
                'filename': Path(image_path).name,
                'apriltag_detected': bool(result.get('apriltag_detected', False)),
                'px_per_mm': float(round(result.get('px_per_mm', 0), 4)),
                'mushrooms': []
            }

            # 处理每个香菇的结果
            mushroom_results = result.get('mushrooms', [])
            for i, mushroom in enumerate(mushroom_results):
                mushroom_data = {
                    'id': int(i + 1),
                    'diameter_mm': float(round(mushroom.get('diameter_mm', 0), 4)),
                    'crack_ratio_percent': float(round(mushroom.get('crack_ratio_percent', 0), 2)),
                    'cap_area_mm2': float(round(mushroom.get('cap_area_mm2', 0), 2)),
                    'confidence': float(round(mushroom.get('detection_confidence', 0), 3)),
                    'pattern_type': str(mushroom.get('pattern_type', '未知')),
                    'pattern_category': str(mushroom.get('pattern_category', '未知')),
                    'white_pattern_ratio': float(round(mushroom.get('white_pattern_ratio', 0) * 100, 2)),
                    'brown_pattern_ratio': float(round(mushroom.get('brown_pattern_ratio', 0) * 100, 2)),
                    'flatness_score': float(round(mushroom.get('flatness_score', 0), 4)),
                    'suspected_ban_gu': bool(mushroom.get('suspected_ban_gu', False))
                }
                response_data['mushrooms'].append(mushroom_data)

            # 编码图像
            response_data['orig_image'] = image_to_base64(orig_image)
            response_data['vis_image'] = image_to_base64(vis_image)

            return jsonify(response_data)

        except Exception as e:
            logging.error(f"分析示例图像出错: {str(e)}", exc_info=True)
            return jsonify({'success': False, 'error': f'服务器错误: {str(e)}'}), 500

    return app


if __name__ == '__main__':
    app = create_app()
    print("=" * 60)
    print("香菇分析Web应用启动")
    print("访问地址: http://localhost:5000")
    print("=" * 60)
    app.run(host='0.0.0.0', port=5000, debug=False)
