# -*- coding: utf-8 -*-
import sqlite3
import json
import logging
import uuid
import os
from typing import Optional, List, Dict
from pathlib import Path
from datetime import datetime
from enum import Enum

logger = logging.getLogger(__name__)

class TaskStatus(Enum):
    PENDING = "pending"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

class ImageTaskManager:
    def __init__(self, db_path: str = None):
        if db_path is None:
            db_path = Path(__file__).parent / "image_tasks.db"
        self.db_path = db_path
        self._init_db()

    def _init_db(self):
        self.db_path.parent.mkdir(parents=True, exist_ok=True)
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        # 核心任务表
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS image_tasks (
                task_id TEXT PRIMARY KEY,
                customer_id TEXT NOT NULL,
                platform TEXT DEFAULT 'qianniu',
                original_image TEXT NOT NULL,
                operation TEXT DEFAULT 'enhance',
                requirements TEXT,
                status TEXT DEFAULT 'pending',
                result_image TEXT,
                price REAL DEFAULT 0.0,
                outcome TEXT DEFAULT 'pending', 
                created_at TEXT,
                updated_at TEXT
            )
        ''')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_status ON image_tasks(status)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_cust_plat ON image_tasks(customer_id, platform)')
        # 兼容旧库：补齐缺失字段
        cursor.execute("PRAGMA table_info(image_tasks)")
        existing_cols = {row[1] for row in cursor.fetchall()}
        if "outcome" not in existing_cols:
            cursor.execute("ALTER TABLE image_tasks ADD COLUMN outcome TEXT DEFAULT 'pending'")
        if "price" not in existing_cols:
            cursor.execute("ALTER TABLE image_tasks ADD COLUMN price REAL DEFAULT 0.0")
        conn.commit()
        conn.close()

    def _get_conn(self):
        conn = sqlite3.connect(self.db_path)
        conn.row_factory = sqlite3.Row
        return conn

    def add_task(self, customer_id: str, platform: str, original_image: str, operation: str, requirements: str = "", status: str = "pending") -> str:
        task_id = str(uuid.uuid4())
        now = datetime.now().isoformat()
        try:
            conn = self._get_conn()
            cursor = conn.cursor()
            cursor.execute('''
                INSERT INTO image_tasks (task_id, customer_id, platform, original_image, operation, requirements, status, created_at, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            ''', (task_id, customer_id, platform, original_image, operation, requirements, status, now, now))
            conn.commit()
            conn.close()
            return task_id
        except Exception as e:
            logger.error(f"Failed to add task: {e}")
            return ""

    def update_status(self, task_id: str, status: str, result_image: str = ""):
        now = datetime.now().isoformat()
        try:
            conn = self._get_conn()
            cursor = conn.cursor()
            if result_image:
                cursor.execute('UPDATE image_tasks SET status = ?, result_image = ?, updated_at = ? WHERE task_id = ?', 
                             (status, result_image, now, task_id))
            else:
                cursor.execute('UPDATE image_tasks SET status = ?, updated_at = ? WHERE task_id = ?', 
                             (status, now, task_id))
            conn.commit()
            conn.close()
        except Exception as e:
            logger.error(f"Failed to update task status: {e}")

    def update_price(self, customer_id: str, platform: str, price: float):
        """记录任务的成交价格"""
        now = datetime.now().isoformat()
        try:
            conn = self._get_conn()
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE image_tasks 
                SET price = ?, updated_at = ? 
                WHERE task_id = (
                    SELECT task_id FROM image_tasks 
                    WHERE customer_id = ? AND platform = ?
                    ORDER BY created_at DESC LIMIT 1
                )
            ''', (price, now, customer_id, platform))
            conn.commit()
            conn.close()
            logger.info(f"[DB] 客户 {customer_id} 任务价格更新为: ¥{price}")
        except Exception as e:
            logger.error(f"Failed to update price: {e}")

    def update_outcome(self, customer_id: str, platform: str, outcome: str):
        """记录任务的最终结局（用于训练样本分类）"""
        now = datetime.now().isoformat()
        try:
            conn = self._get_conn()
            cursor = conn.cursor()
            cursor.execute('''
                UPDATE image_tasks 
                SET outcome = ?, updated_at = ? 
                WHERE task_id = (
                    SELECT task_id FROM image_tasks 
                    WHERE customer_id = ? AND platform = ?
                    ORDER BY created_at DESC LIMIT 1
                )
            ''', (outcome, now, customer_id, platform))
            conn.commit()
            conn.close()
            logger.info(f"[DB] 客户 {customer_id} 任务结局更新为: {outcome}")
        except Exception as e:
            logger.error(f"Failed to update outcome: {e}")

# 单例
db = ImageTaskManager()