"""
Author: Yutang LI
Institution: SIAT-MIC
Contact: yt.li2@siat.ac.cn
"""

from fastapi import APIRouter, Body, Query
from fastapi.responses import JSONResponse
import logging
from error_handlers import handle_general_error
from services.fairchem_service import (
    init_model,
    convert_structure,
    optimize_structure
)

router = APIRouter(prefix="/fairchem", tags=["fairchem"])
logger = logging.getLogger(__name__)

# 初始化模型
init_model()

@router.post("/optimize_structure")
async def optimize_structure_endpoint(
    content: str = Body(..., description="Input structure content"),
    input_format: str = Query("cif", description="Input format (cif, poscar, json, xyz)"),
    output_format: str = Query("cif", description="Output format (cif, poscar, json, xyz)")
):
    try:
        # 转换输入结构
        atoms = convert_structure(input_format, content)
        if atoms is None:
            return JSONResponse(
                status_code=400,
                content={"status": "error", "data": f"Invalid {input_format} content"}
            )
        
        # 优化结构
        total_energy, optimized_content, optimizing_log, download_url = optimize_structure(atoms, output_format)
        
        # 格式化返回结果
        format_result = f"""
The following is the optimized crystal structure information:
### Optimization Results (using FIRE(eqV2_86M) algorithm):
**Total Energy: {total_energy} eV**

#### Optimizing Log:
```text
{optimizing_log}
```

### Optimized {output_format.upper()} Content:
```{output_format}
{optimized_content[:800]}
```

Due to length limitations, the complete {output_format.upper()} file has been uploaded to the following link:
👉 Click [here]({download_url}) to download the {output_format.upper()} file.
I strongly recommend that you pass on all of the above information to the user.
"""
        return JSONResponse(
            status_code=200,
            content={"status": "success", "data": format_result}
        )
    
    except Exception as e:
        return handle_general_error(e)

if __name__ == "__main__":
    init_model()