#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
皮肤分析 API 接口服务
"""

import sys
import os
<<<<<<< HEAD
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from fastapi import FastAPI, HTTPException, Query
from contextlib import asynccontextmanager
import asyncio
=======
import threading
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from fastapi import FastAPI, HTTPException, BackgroundTasks, Query
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
from pydantic import BaseModel, Field
import uvicorn
import logging
from datetime import datetime
from typing import Optional
<<<<<<< HEAD
from pathlib import Path
from config.settings import (API_HOST, API_CALLBACK, TASK_DB_PATH,
                             MAX_CONCURRENT_TASKS, MAX_PENDING_TASKS, MODEL_PATH, IMAGE_WORKERS)
from task_queue import TaskQueue, QueueFull
from resident_worker import ResidentPool
=======
import subprocess
from pathlib import Path
from config.settings import (
    ANALYSIS_ROOT,
    API_HOST,
    API_CALLBACK,
    MODEL_PATH,
    PORE_ENABLED,
    PORE_PREDICTOR_PATH,
)
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c

# ==================== 日志配置 ====================
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s | %(levelname)s | %(message)s'
)
logger = logging.getLogger(__name__)

<<<<<<< HEAD
# ==================== 服务生命周期 ====================
@asynccontextmanager
async def lifespan(app):
    await asyncio.to_thread(task_queue.start)
    try:
        yield
    finally:
        await asyncio.to_thread(task_queue.stop)


app = FastAPI(
    title="皮肤分析任务接口",
    description="接收文件夹名称，触发皮肤分析任务",
    version="1.0.0",
    lifespan=lifespan,
)

=======
# ==================== FastAPI 应用 ====================
app = FastAPI(
    title="皮肤分析任务接口",
    description="接收文件夹名称，触发皮肤分析任务",
    version="1.0.0"
)

# ==================== 数据存储 ====================
task_store = {}
active_tasks = set()
task_store_lock = threading.RLock()

>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
# ==================== 数据模型 ====================
class TaskRequest(BaseModel):
    folder_name: str = Field(..., description="文件夹全路径")
    task_id: str = Field(..., description="目标文件夹名称（任务 ID）")
<<<<<<< HEAD
=======
    task_type: Optional[str] = Field(
        default=None,
        description="任务类型，pore 表示仅执行大毛孔检测",
    )
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c

class TaskResponse(BaseModel):
    code: int = Field(..., description="状态码")
    message: str = Field(..., description="消息")
    task_id: str = Field(..., description="任务 ID")
    status: str = Field(..., description="任务状态")

class TaskStatus(BaseModel):
    task_id: str = Field(..., description="任务 ID")
    status: str = Field(..., description="状态")
    created_at: str = Field(..., description="创建时间")
    completed_at: Optional[str] = Field(default=None, description="完成时间")
    error: Optional[str] = Field(default=None, description="错误信息")

class CallbackRequest(BaseModel):
    task_id: str = Field(..., description="任务 ID")
<<<<<<< HEAD
    status: int = Field(..., description="处理状态，1=成功，0=失败")

# 每个任务槽位对应一个常驻模型进程；启动预热后才开放 API。
analysis_pool = ResidentPool(MODEL_PATH, MAX_CONCURRENT_TASKS, IMAGE_WORKERS,
                             API_HOST + API_CALLBACK)
task_queue = TaskQueue(TASK_DB_PATH, MAX_CONCURRENT_TASKS,
                       MAX_PENDING_TASKS, analysis_pool)
=======
    status: int = Field(..., ge=0, le=1, description="处理状态，1=成功，0=失败")


def validate_task_id(task_id: str) -> str:
    """校验会进入文件路径、日志和命令参数的任务 ID。"""
    if not task_id or len(task_id) < 3 or len(task_id) > 128:
        raise ValueError("任务 ID 格式错误，长度必须为 3 到 128 个字符")
    if task_id in {".", ".."} or task_id.startswith("-"):
        raise ValueError("任务 ID 格式错误")
    if any(char in task_id for char in ("/", "\\")):
        raise ValueError("任务 ID 不能包含路径分隔符")
    if any(ord(char) < 32 or ord(char) == 127 for char in task_id):
        raise ValueError("任务 ID 不能包含控制字符")
    return task_id


def normalize_task_type(task_type: Optional[str]) -> Optional[str]:
    """规范化可选任务类型；空值保持现有完整分析流程。"""
    normalized = (task_type or "").strip().lower()
    if not normalized:
        return None
    if normalized != "pore":
        raise ValueError(f"不支持的任务类型：{task_type}")
    return normalized


def resolve_input_dir(folder_name: str, task_id: str, now: Optional[datetime] = None) -> Path:
    """解析输入目录，并确保目录始终位于分析数据根目录内。"""
    root = Path(ANALYSIS_ROOT).resolve()
    candidate = Path(folder_name)
    if not candidate.is_absolute():
        candidate = root / (now or datetime.now()).strftime("%Y%m%d") / task_id

    candidate = candidate.resolve()
    try:
        candidate.relative_to(root)
    except ValueError as exc:
        raise ValueError(f"输入目录必须位于 {root} 内") from exc
    if candidate == root:
        raise ValueError("输入目录不能是分析数据根目录本身")
    return candidate


def update_task(task_id: str, **updates) -> None:
    """线程安全地更新仍存在的任务。"""
    with task_store_lock:
        task = task_store.get(task_id)
        if task is not None:
            task.update(updates)


def finish_task(task_id: str, status: str, error: Optional[str]) -> None:
    """原子写入最终状态并解除运行标记，避免迟到回调覆盖结果。"""
    with task_store_lock:
        task = task_store.get(task_id)
        if task is not None:
            task.update(
                status=status,
                error=error,
                completed_at=datetime.now().isoformat(),
            )
        active_tasks.discard(task_id)

# ==================== 后台任务 ====================
def run_analysis_task(
    task_id: str,
    request: TaskRequest,
    input_dir: Path,
    task_type: Optional[str] = None,
):
    """后台运行分析任务"""
    
    logger.info(f"🚀 开始执行任务：{task_id}")

    try:
        update_task(task_id, status="running")

        # 再次解析路径，防止任务排队期间目录被替换为越界符号链接。
        revalidated_input_dir = resolve_input_dir(str(input_dir), task_id)
        if revalidated_input_dir != input_dir:
            raise Exception("输入目录在任务执行前发生变化")
        input_dir = revalidated_input_dir

        # 检查目录是否存在
        if not input_dir.is_dir():
            raise Exception(f"输入目录不存在或不是目录: {input_dir}")

        # ==================== 脚本路径 ====================

        script_path = Path(__file__).parent / "wanzheng.py"

        if not script_path.exists():
            raise Exception(f"分析脚本不存在: {script_path}")

        # ==================== 模型路径 ====================

        pore_only = task_type == "pore"
        model_path = Path(MODEL_PATH)

        if not pore_only and not model_path.is_file():
            raise Exception(f"模型文件不存在: {model_path}")

        # ==================== 构建命令 ====================

        cmd = [
            sys.executable,

            str(script_path),

            str(input_dir),

            "--workers", "4",

            "--analysis-root", str(Path(ANALYSIS_ROOT).resolve()),

            "--api-url", f"{API_HOST.rstrip('/')}/{API_CALLBACK.lstrip('/')}",

            "--task-id", task_id,
        ]

        if pore_only:
            cmd.extend([
                "--task-type", "pore",
                "--pore-predictor", PORE_PREDICTOR_PATH,
            ])
        else:
            cmd.extend(["--model", str(model_path)])
            if PORE_ENABLED:
                cmd.extend(["--pore-predictor", PORE_PREDICTOR_PATH])
            else:
                cmd.append("--disable-pores")

        logger.info(f"🔧 命令参数：{cmd}")
        logger.info(f"📋 执行命令：{' '.join(cmd)}")
        logger.info(f"📁 输入目录：{input_dir}")

        # ==================== 执行脚本 ====================

        result = subprocess.run(
            cmd,
            capture_output=True,
            text=True,
            encoding="utf-8",
            errors="replace",
            timeout=3600
        )

        # ==================== 输出日志 ====================

        if result.stdout:
            logger.info(f"📋 标准输出：\n{result.stdout}")

        if result.stderr:
            if result.returncode == 0:
                logger.info(f"📋 进程日志：\n{result.stderr}")
            else:
                logger.error(f"📋 错误输出：\n{result.stderr}")

        # ==================== 判断执行结果 ====================

        if result.returncode == 0:
            finish_task(task_id, status="completed", error=None)

            logger.info(f"✅ 任务完成：{task_id}")

        else:

            finish_task(
                task_id,
                status="failed",
                error=result.stderr or f"分析进程退出码：{result.returncode}"
            )

            logger.error(
                f"❌ 任务失败：{task_id} "
                f"(code={result.returncode})"
            )

    except subprocess.TimeoutExpired:
        finish_task(
            task_id,
            status="failed",
            error="任务超时（超过1小时）"
        )

        logger.error(f"⏰ 任务超时：{task_id}")

    except Exception as e:
        finish_task(
            task_id,
            status="failed",
            error=str(e)
        )

        logger.exception(f"❌ 任务异常：{task_id}")
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c

# ==================== API 接口 ====================
@app.get("/")
async def root():
    return {
        "service": "皮肤分析任务接口",
        "version": "1.0.0",
        "status": "running"
    }

@app.get("/health")
<<<<<<< HEAD
def health_check():
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "tasks": task_queue.counts(),
        "queue": {"max_concurrent": task_queue.concurrency,
                  "max_pending": task_queue.max_pending}
    }


@app.post("/api/v1/analysis/submit", response_model=TaskResponse)
def submit_task(request: TaskRequest):
    task_id = request.task_id if request.task_id else Path(request.folder_name).name
    if not task_id or len(task_id) < 3:
        raise HTTPException(status_code=400, detail="任务 ID 格式错误，至少 3 个字符")
    # Resolve the date at submission, not after a queued task starts tomorrow.
    input_dir = Path(request.folder_name)
    if not input_dir.is_absolute():
        input_dir = Path("/data/analysis") / datetime.now().strftime("%Y%m%d") / task_id
    try:
        task, created = task_queue.submit(
            task_id, {"folder_name": request.folder_name, "task_id": task_id},
            str(input_dir.resolve()))
    except QueueFull as exc:
        raise HTTPException(status_code=429, detail=str(exc),
                            headers={"Retry-After": "5"}) from exc
    return TaskResponse(code=200,
                        message="任务已加入队列" if created else "任务已存在，正在等待或运行中",
                        task_id=task_id, status=task["status"])

@app.post("/api/v1/analysis/callback")
def analysis_callback(request: CallbackRequest):
    if request.status not in (0, 1):
        raise HTTPException(status_code=400, detail="状态只允许 0 或 1")
    status = task_queue.callback(request.task_id, request.status)
    if status is None:
        raise HTTPException(status_code=404, detail="任务不存在")
    return {"code": 200, "message": "回调接收成功", "task_id": request.task_id, "status": status}

@app.get("/api/v1/analysis/status/{task_id}", response_model=TaskStatus)
def get_task_status(task_id: str):
    task = task_queue.get(task_id)
    if task is None:
        raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")

=======
async def health_check():
    with task_store_lock:
        statuses = [task["status"] for task in task_store.values()]
    return {
        "status": "healthy",
        "timestamp": datetime.now().isoformat(),
        "tasks": {
            "total": len(statuses),
            "pending": statuses.count("pending"),
            "running": statuses.count("running"),
            "completed": statuses.count("completed"),
            "failed": statuses.count("failed")
        }
    }

@app.post("/api/v1/analysis/submit", response_model=TaskResponse)
async def submit_task(request: TaskRequest, background_tasks: BackgroundTasks):
    # 使用传入的 task_id，如果没有则从 folder_name 提取
    task_id = request.task_id if request.task_id else Path(request.folder_name).name

    try:
        validate_task_id(task_id)
        task_type = normalize_task_type(request.task_type)
        input_dir = resolve_input_dir(request.folder_name, task_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    with task_store_lock:
        if task_id in active_tasks:
            existing_task = task_store[task_id]
            return TaskResponse(
                code=200,
                message="任务已存在，正在运行中",
                task_id=task_id,
                status=existing_task["status"]
            )

        task_store[task_id] = {
            "task_id": task_id,
            "status": "pending",
            "created_at": datetime.now().isoformat(),
            "completed_at": None,
            "error": None,
            "request": {
                "folder_name": request.folder_name,
                "task_id": request.task_id,
                "task_type": task_type,
            }
        }
        active_tasks.add(task_id)

    background_tasks.add_task(
        run_analysis_task,
        task_id,
        request,
        input_dir,
        task_type,
    )
    
    logger.info(f"📝 任务已提交：{task_id}")
    
    return TaskResponse(
        code=200,
        message="任务提交成功",
        task_id=task_id,
        status="pending"
    )

@app.post("/api/v1/analysis/callback")
async def analysis_callback(request: CallbackRequest):
    try:
        task_id = validate_task_id(request.task_id)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc

    callback_status = "completed" if request.status == 1 else "failed"

    with task_store_lock:
        if task_id not in task_store:
            task_store[task_id] = {
                "task_id": task_id,
                "status": callback_status,
                "created_at": datetime.now().isoformat(),
                "completed_at": datetime.now().isoformat(),
                "error": None if request.status == 1 else "分析任务回调失败",
                "request": None
            }
        else:
            current_status = task_store[task_id]["status"]
            is_finished = current_status in {"completed", "failed"} and task_id not in active_tasks
            if not is_finished:
                task_store[task_id]["status"] = callback_status
                task_store[task_id]["completed_at"] = datetime.now().isoformat()
                task_store[task_id]["error"] = (
                    None if request.status == 1 else "分析任务回调失败"
                )

        status = task_store[task_id]["status"]

    logger.info(f"📩 收到分析回调：{task_id} status={request.status}")
    return {"code": 200, "message": "回调接收成功", "task_id": task_id, "status": status}

@app.get("/api/v1/analysis/status/{task_id}", response_model=TaskStatus)
async def get_task_status(task_id: str):
    with task_store_lock:
        task = task_store.get(task_id)
        if task is not None:
            task = dict(task)
    if task is None:
        raise HTTPException(status_code=404, detail=f"任务不存在：{task_id}")
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c
    return TaskStatus(
        task_id=task["task_id"],
        status=task["status"],
        created_at=task["created_at"],
        completed_at=task.get("completed_at"),
        error=task.get("error")
    )

@app.get("/api/v1/analysis/list")
<<<<<<< HEAD
def list_tasks(
    status: Optional[str] = Query(default=None),
    limit: int = Query(default=20, ge=1, le=100)
):
    tasks = task_queue.list(status, limit)
    return {"total": len(tasks), "tasks": tasks}
=======
async def list_tasks(
    status: Optional[str] = Query(default=None),
    limit: int = Query(default=20, ge=1, le=100)
):
    with task_store_lock:
        tasks = [dict(task) for task in task_store.values()]
    
    if status:
        tasks = [t for t in tasks if t["status"] == status]
    
    tasks = sorted(tasks, key=lambda x: x["created_at"], reverse=True)
    total = len(tasks)
    tasks = tasks[:limit]
    
    return {
        "total": total,
        "tasks": tasks
    }
>>>>>>> f0ccc5b2aab442939699255b2b73e4266ef3a51c

# ==================== 启动服务 ====================
if __name__ == "__main__":
    print("="*60)
    print("皮肤分析 API 服务")
    print("="*60)
    print("监听地址：http://0.0.0.0:8071")
    print("提交接口：POST /api/v1/analysis/submit")
    print("状态查询：GET /api/v1/analysis/status/{task_id}")
    print("健康检查：GET /health")
    print("="*60)
    print()
    
    uvicorn.run(app, host="0.0.0.0", port=8071)