from fastapi import FastAPI, File, UploadFile, HTTPException, Request
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
import pytesseract
from PIL import Image
import io
import os
import logging

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(
    title="OCR Text Recognition API",
    description="Extract text from images using Tesseract OCR",
    version="1.0.0"
)

# CORS 配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# API Key 验证（可选）
API_KEY = os.environ.get("API_KEY", "")

def verify_api_key(request: Request):
    """验证 API Key"""
    if not API_KEY:
        return True  # 未设置 API Key，允许所有请求
    
    auth_header = request.headers.get("Authorization", "")
    if auth_header.startswith("Bearer "):
        token = auth_header[7:]
        if token == API_KEY:
            return True
    
    x_api_key = request.headers.get("X-API-Key", "")
    if x_api_key == API_KEY:
        return True
    
    return False

@app.post("/ocr")
async def extract_text(
    request: Request,
    file: UploadFile = File(...),
    lang: str = "eng"
):
    """
    从图片中提取文字
    
    参数：
    - file: 图片文件（支持 PNG, JPG, JPEG, WEBP, BMP, TIFF）
    - lang: 语言代码（默认 eng，可选 chi_sim 中文简体、chi_tra 中文繁体、jpn 日语等）
    
    返回：
    - text: 识别出的文字
    - confidence: 置信度（0-100）
    """
    # API Key 验证
    if not verify_api_key(request):
        # 免费用户 IP 限制
        client_ip = request.headers.get("X-Forwarded-For", request.client.host)
        logger.warning(f"Unauthorized access from {client_ip}")
        raise HTTPException(status_code=401, detail="Invalid API Key. Please subscribe to use this API.")
    
    # 检查文件类型
    allowed_types = ["image/png", "image/jpeg", "image/jpg", "image/webp", "image/bmp", "image/tiff"]
    if file.content_type not in allowed_types:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file type: {file.content_type}. Supported: PNG, JPG, JPEG, WEBP, BMP, TIFF"
        )
    
    try:
        # 读取图片
        content = await file.read()
        image = Image.open(io.BytesIO(content))
        
        # 转换为 RGB（如果需要）
        if image.mode in ("RGBA", "P"):
            image = image.convert("RGB")
        
        # OCR 识别
        data = pytesseract.image_to_data(image, lang=lang, output_type=pytesseract.Output.DICT)
        
        # 提取文字和置信度
        text_parts = []
        confidences = []
        
        for i, text in enumerate(data["text"]):
            if text.strip():  # 只保留非空文字
                text_parts.append(text)
                confidences.append(data["conf"][i])
        
        # 合并文字
        full_text = " ".join(text_parts)
        
        # 计算平均置信度
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        
        return JSONResponse({
            "success": True,
            "text": full_text,
            "confidence": round(avg_confidence, 2),
            "language": lang,
            "word_count": len(text_parts)
        })
        
    except pytesseract.TesseractError as e:
        logger.error(f"Tesseract error: {str(e)}")
        raise HTTPException(status_code=400, detail=f"OCR Error: {str(e)}")
    except Exception as e:
        logger.error(f"Error processing image: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error processing image: {str(e)}")

@app.post("/ocr/detailed")
async def extract_text_detailed(
    request: Request,
    file: UploadFile = File(...),
    lang: str = "eng"
):
    """
    从图片中提取文字（详细版，包含每个文字的位置信息）
    
    参数：
    - file: 图片文件
    - lang: 语言代码
    
    返回：
    - text: 完整文字
    - words: 每个词的详细信息（文字、位置、置信度）
    - confidence: 平均置信度
    """
    # API Key 验证
    if not verify_api_key(request):
        raise HTTPException(status_code=401, detail="Invalid API Key. Please subscribe to use this API.")
    
    # 检查文件类型
    allowed_types = ["image/png", "image/jpeg", "image/jpg", "image/webp", "image/bmp", "image/tiff"]
    if file.content_type not in allowed_types:
        raise HTTPException(
            status_code=400, 
            detail=f"Unsupported file type: {file.content_type}"
        )
    
    try:
        # 读取图片
        content = await file.read()
        image = Image.open(io.BytesIO(content))
        
        if image.mode in ("RGBA", "P"):
            image = image.convert("RGB")
        
        # OCR 识别
        data = pytesseract.image_to_data(image, lang=lang, output_type=pytesseract.Output.DICT)
        
        # 提取详细文字信息
        words = []
        confidences = []
        
        for i, text in enumerate(data["text"]):
            if text.strip():
                words.append({
                    "text": text,
                    "confidence": data["conf"][i],
                    "position": {
                        "x": data["left"][i],
                        "y": data["top"][i],
                        "width": data["width"][i],
                        "height": data["height"][i]
                    }
                })
                confidences.append(data["conf"][i])
        
        full_text = " ".join([w["text"] for w in words])
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        
        return JSONResponse({
            "success": True,
            "text": full_text,
            "words": words,
            "confidence": round(avg_confidence, 2),
            "language": lang,
            "word_count": len(words)
        })
        
    except Exception as e:
        logger.error(f"Error: {str(e)}")
        raise HTTPException(status_code=500, detail=f"Error: {str(e)}")

@app.get("/")
async def root():
    """API 信息"""
    return {
        "name": "OCR Text Recognition API",
        "version": "1.0.0",
        "description": "Extract text from images using Tesseract OCR",
        "endpoints": {
            "/ocr": "Extract text from image (POST)",
            "/ocr/detailed": "Extract text with position details (POST)"
        },
        "supported_languages": {
            "eng": "English",
            "chi_sim": "Chinese Simplified",
            "chi_tra": "Chinese Traditional",
            "jpn": "Japanese",
            "kor": "Korean",
            "fra": "French",
            "deu": "German",
            "spa": "Spanish",
            "rus": "Russian",
            "ara": "Arabic"
        }
    }

@app.get("/health")
async def health():
    """健康检查"""
    return {"status": "ok"}