from mcp.server.fastmcp import FastMCP
from .api.imagex import ImagexAPI
import uuid
import json
import os


def Error(message: str):
    return "API Error: " + message


def HandlerVolcResponse(response: dict):
    if (
        response
        and hasattr(response, "ResponseMetadata")
        and response.ResponseMetadata
        and hasattr(response.ResponseMetadata, "Error")
        and response.ResponseMetadata.Error
    ):
        return Error(response.ResponseMetadata.Error.Message)
    return str(response)


def create_mcp_server():
    mcp = FastMCP(
        "VeImageX MCP",
        description="Volcengine(火山引擎) ImageX(图片服务) MCP , 你的图片处理存储分发助手",
    )
    imagex_service = ImagexAPI()
    service_id = imagex_service.service_id
    domain = imagex_service.domain

    @mcp.tool()
    def guide():
        """VeImageX MCP is the Volcengine(火山引擎) ImageX(图片服务) MCP Server, an image processing assistant that helps users upload images, process images, distribute images, and obtain image URLs.

        Before using VeImageX MCP, you need to know a few things:
        1. For tools that rely on service_id and domain, prioritize obtaining from environment variable as the default value. If no value is specified in the environment variable, prompt the user.

        2. For tools that rely on template, the default value of the template is tplv-service_id-image, where service_id needs to be replaced with the actual service_id, e.g., if the user's service_id is n9b2vwdhz3, use tplv-n9b2vwdhz3-image
        
        3. `upload_image` is used to upload images, using local absolute paths - the tool will help you upload to Volcengine using fs. If you don't have this path, you try to find user's intent, or ask him.

        4. `get_image_url_by_store_uri` is used to get the image URL
        
        5. If the result generated by the tool is a URI, automatically call `get_image_url_by_store_uri` to get an accessible URL and return it.
        """
        return """use `guide` description to get how to use ImageX Mcp Server"""

    @mcp.tool()
    def get_all_services_resource() -> str:
        """Get list of all service information.
        Important information includes:
        1. Domain Name - DomainName
        2. Service ID - ServiceId
        3. Service Name - ServiceName
        4. Service Type - ServiceType: StaticRc (static hosting), Image (image processing)
        5. Service Status - ServiceStatus"""
        result = imagex_service.get_all_image_services({})
        return str(HandlerVolcResponse(result))

    @mcp.tool()
    def upload_image(file_path: list[str]) -> str:
        """Upload images to specified service.

        Args:
            file_path: List of image file paths, use absolute path."""
        if not service_id:
            return Error(
                "service_id is required, please use get_all_services_resource to get the service_id"
            )
        if not file_path:
            return Error("file_path is required")
        params = {"ServiceId": service_id, "SkipMeta": False, "SkipCommit": False}

        result = imagex_service.upload_image(params, file_path)

        return str(HandlerVolcResponse(result))

    @mcp.tool()
    def get_all_image_templates() -> str:
        """Get all templates (tpl) for the Image service. Call this before getting image URLs to query the corresponding API service.
        1. If search is empty, get all templates"""

        if not service_id:
            return Error(
                "service_id is required, please use get_all_services_resource to get the service_id"
            )
        result = imagex_service.get_all_image_templates(
            {"ServiceId": service_id, "Limit": 1000}
        )

        return str(HandlerVolcResponse(result))

    @mcp.tool()
    def get_image_url_by_store_uri(
        uri: str = None
    ) -> str:
        """
        Get image access URL through URI and template information. 
        URI is obtained from upload_image.

        Args:
            uri: The storage URI of the image
        """
        
        if not service_id:
            return Error(
                "service_id is required, please use get_all_services_resource to get the service_id"
            )
        if not domain:
            return Error("domain is required")
        if not uri:
            return Error("uri is required")

        return str("https://{domain}/{uri}~tplv-{serviceid}-image.image".format(
            domain = domain,
            uri = uri,
            serviceid = service_id
        ))

    @mcp.tool()
    def generate_image_by_text(
        text: str = None
    ) -> str:
        """Generate an image based on the text.

        Args:
            text: The text to generate an image from
        """
        
        if not service_id:
            return Error(
                "service_id is required, please use get_all_services_resource to get the service_id"
            )
        if not domain:
            return Error("domain is required")
        
        tpl = "tplv-{serviceid}-image.image".format(
            serviceid = service_id
        )
        
        if not tpl:
            return Error("tpl is required")
        if not text:
            return Error("text is required")

        query = {
            "ServiceId": service_id,
        }

        params = {
            "Domain": domain,
            "Template": tpl,
            "Overwrite": True,
            "ReqJson": {"req_key": "high_aes_general_v20_L", "prompt": text},
            "ModelAction": "CVProcess",
            "ModelVersion": "2022-08-31",
            "Outputs": [f"output_{uuid.uuid4().hex[:8]}"],
        }

        result = imagex_service.get_cv_text_generate_image(query, params)
        return str(HandlerVolcResponse(result))
    
    @mcp.tool()
    def enhance_image_quality(
        input_image_url: str,
    ) -> str:
        """Under the premise of ensuring consistency with the original image content, significantly improve the clarity.

        Args:
            input_image_url: The input image url
        """
        
        workflow_parameter = {
            "Input": {
                "ObjectKey": input_image_url,
                "DataType": "url",
            },
            "GenDREnhanceParam": {
                "Multiple": 2,  
            },
        }

        params = {
            "ServiceId": service_id,
            "WorkflowTemplateId": "system_workflow_ai_super_resolution",
            "WorkflowParameter": json.dumps(workflow_parameter),
        }

        result = imagex_service.ai_process({}, params)
        
        return str(HandlerVolcResponse(result))
    
    @mcp.tool()
    def convert_image_to_comic_style(
        input_image_url: str,
    ) -> str:
        """ Turn the input image into a comic style.

        Args:
            input_image_url: The input image url
        """

        params = {
            "ServiceId": service_id,
            "StoreUri": input_image_url,
        }

        result = imagex_service.get_image_comic_result(params)
        
        return str(HandlerVolcResponse(result))
    
    @mcp.tool()
    def image_ocr(
        input_image_url: str,
    ) -> str:
        """ Recognize the text in the image.

        Args:
            input_image_url: The input image url
        """

        query = {
            "ServiceId": service_id,
        }
        params = {
            "ImageUrl": input_image_url,
            "Scene": "general",
        }

        result = imagex_service.get_image_ocr_v2(query, params)
        
        return str(HandlerVolcResponse(result))
    
    @mcp.tool()
    def expand_image(
        input_image_url: str,
    ) -> str:
        """ Expand the image to the surrounding area.

        Args:
            input_image_url: The input image url
        """

        params = {
            "ServiceId": service_id,
            "StoreUri": input_image_url,
            "Model": 1,
            "Top": 0.2,
            "Bottom": 0.2,
            "Left": 0.2,
            "Right": 0.2,
        }

        result = imagex_service.get_image_bg_fill_result(params)
        
        return str(HandlerVolcResponse(result))
    
    @mcp.tool()
    def evaluate_image_quality(
        input_image_url: str,
    ) -> str:
        """ Evaluate the image quality. There will be NrScoreResult and FrScoreResult fields in the returned results.
            Summary of NrScoreResult and FrScoreResult Fields
            1. NrScoreResult (No-Reference Image Quality Assessment)
            Purpose: Evaluates absolute image quality without requiring a reference image.
            Sub-fields:
            VqScore (0-100): Overall perceptual quality score (higher = better)
            Noise (0-100): Noise level (higher = more noise)
            Blur (0-100): Blurriness (higher = more blurred)
            Blockiness (0-100): Compression blocking artifacts
            Brightness (0-100): Image luminance level
            Contrast (0-100): Color contrast intensity
            AdvColor (0-100): Advanced color quality assessment
            Aesthetic (0-100): Aesthetic appeal score
            FaceDetect (0/1): Whether faces are detected
            FaceScore (0-100): Quality score of detected faces
            FaceNum: Number of faces detected
            TextDetect (0/1): Whether text is detected
            TextScore (0-100): Quality score of detected text
            TextNum: Number of text regions detected
            Use Case: Standalone quality assessment for single images.
            
            2. FrScoreResult (Full-Reference Image Quality Assessment)
            Purpose: Compares target image against a reference image.
            Sub-fields:
            Psnr (0-60): Peak Signal-to-Noise Ratio (higher = better)
            Ssim (0-1): Structural Similarity Index
            Vmaf (0-100): Video Multi-method Assessment Fusion score
            DeltaE (0-100): Color difference metric (lower = better)
            Fid (0-100): Feature similarity index
            Nab (0-100): Noise and blur difference score
            Requirements:
            Requires reference image (ImageUrlRef)
            Both images must have identical resolution
            Use Case: Quality comparison between original and processed images.

        Args:
            input_image_url: The input image url
        """
        
        query = {
            "ServiceId": service_id,
        }

        params = {
            "ImageUrl": input_image_url,
            "VqType": "nr_index,vqscore,advcolor,blockiness,noise,aesmod,blur,cg,contrast,texture,brightness,overexposure,hue,saturation,green,cmartifacts", 
        }

        result = imagex_service.get_image_quality(query, params)
        
        return str(HandlerVolcResponse(result))
    
    return mcp
