# -*- coding: utf-8 -*-
# Copyright (c) 2023 Salvador E. Tropea
# Copyright (c) 2023 Instituto Nacional de Tecnología Industrial
# Copyright (c) 2022 uPesy Electronics (Alexis)
# License: AGPL-3.0
# Project: KiBot (formerly KiPlot)
# Adapted from: https://github.com/uPesy/easyeda2kicad.py

from dataclasses import dataclass
import json
import os
import requests
import re
import textwrap
from ..misc import W_EEDA3D
from .. import log

logger = log.get_logger()
API_ENDPOINT = "https://easyeda.com/api/products/{lcsc_id}/components?version=6.4.19.5"
ENDPOINT_3D_MODEL = "https://easyeda.com/analyzer/api/3dmodel/{uuid}"
VRML_HEADER = "#VRML V2.0 utf8\n# 3D model generated by KiBot (using easyeda2kicad.py code)\n"
USER_AGENT = "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:102.0) Gecko/20100101 Firefox/102.0"
# From https://github.com/TousstNicolas/JLC2KiCad_lib/blob/master/JLC2KiCadLib/footprint/model3d.py
# `qAxj6KHrDKw4blvCG8QJPs7Y` is a constant in
# https://modules.lceda.cn/smt-gl-engine/0.8.22.6032922c/smt-gl-engine.js
# and points to the bucket containing the step files.
ENDPOINT_STEP = "https://modules.easyeda.com/qAxj6KHrDKw4blvCG8QJPs7Y/{uuid}"
# KiCad materials that we try to detect
MATERIAL_PIN_01 = {"name": 'PIN-01', "ambient_intensity": 0.271, "shininess": 0.7, "transparency": 0,
                   "diffuse_color": ('0.824', '0.820', '0.781'), "specular_color": ('0.328', '0.258', '0.172')}
MATERIAL_PIN_02 = {"name": 'PIN-02', "ambient_intensity": 0.379, "shininess": 0.4, "transparency": 0,
                   "diffuse_color": ('0.859', '0.738', '0.496'), "specular_color": ('0.137', '0.145', '0.184')}
MATERIAL_MET_01 = {"name": 'MET-01', "ambient_intensity": 0.250, "shininess": 0.056, "transparency": 0,
                   "diffuse_color": ('0.298', '0.298', '0.298'), "specular_color": ('0.398', '0.398', '0.398')}
MATERIAL_EPOXY_04 = {"name": 'IC-BODY-EPOXY-04', "ambient_intensity": 0.293, "shininess": 0.35, "transparency": 0,
                     "diffuse_color": ('0.148', '0.145', '0.145'), "specular_color": ('0.180', '0.168', '0.160')}
if 'KIBOT_EASYEDA_API' in os.environ:
    API_ENDPOINT = os.environ['KIBOT_EASYEDA_API']
if 'KIBOT_EASYEDA_MODEL' in os.environ:
    ENDPOINT_3D_MODEL = os.environ['KIBOT_EASYEDA_MODEL']
if 'KIBOT_EASYEDA_STEP' in os.environ:
    ENDPOINT_STEP = os.environ['KIBOT_EASYEDA_STEP']


class EasyedaApi:
    def __init__(self):
        self.headers = {"Accept-Encoding": "gzip, deflate",
                        "Accept": "application/json, text/javascript, */*; q=0.01",
                        "Content-Type": "application/x-www-form-urlencoded; charset=UTF-8",
                        "User-Agent": USER_AGENT}

    def get_info_from_easyeda_api(self, lcsc_id):
        r = requests.get(url=API_ENDPOINT.format(lcsc_id=lcsc_id), headers=self.headers)
        api_response = r.json()

        if not api_response or "code" in api_response and api_response["success"] is False:
            logger.warning(W_EEDA3D+"Failed EasyEDA request for `{}`, got: {}".format(lcsc_id, api_response))
            return {}

        return r.json()

    def get_cad_data_of_component(self, lcsc_id):
        cp_cad_info = self.get_info_from_easyeda_api(lcsc_id=lcsc_id)
        if cp_cad_info == {}:
            return {}
        return cp_cad_info["result"]

    def get_raw_3d_model_obj(self, uuid):
        # Surface model
        url = ENDPOINT_3D_MODEL.format(uuid=uuid)
        logger.debugl(3, f"- Downloading raw 3D model from {url}")
        r = requests.get(url=url, headers={"User-Agent": self.headers["User-Agent"]})
        if r.status_code != requests.codes.ok:
            logger.warning(W_EEDA3D+"Failed to download 3D model data found for EasyEDA uuid: "+uuid)
            obj = None
        else:
            obj = r.content.decode()
        # 3D object
        url = ENDPOINT_STEP.format(uuid=uuid)
        logger.debugl(3, f"- Downloading STEP 3D model from {url}")
        r = requests.get(url=url, headers={"User-Agent": self.headers["User-Agent"]})
        if r.status_code != requests.codes.ok:
            logger.warning(W_EEDA3D+"Failed to download STEP 3D model data found for EasyEDA uuid: "+uuid)
            step = None
        else:
            step = r.content
        return obj, step


def convert_to_mm(dim: float):
    return float(dim) * 10 * 0.0254


@dataclass
class Ee3dModelBase:
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0

    def convert_to_mm(self) -> None:
        self.x = convert_to_mm(self.x)
        self.y = convert_to_mm(self.y)
        self.z = convert_to_mm(self.z)


@dataclass
class Ee3dModel:
    name: str
    uuid: str
    translation: Ee3dModelBase
    rotation: Ee3dModelBase
    raw_obj: str = None
    step: str = None

    def convert_to_mm(self) -> None:
        self.translation.convert_to_mm()
        # self.translation.z = self.translation.z


@dataclass
class Ki3dModelBase:
    x: float = 0.0
    y: float = 0.0
    z: float = 0.0


@dataclass
class Ki3dModel:
    name: str
    translation: Ki3dModelBase
    rotation: Ki3dModelBase
    raw_wrl: str = None


class Easyeda3dModelImporter:
    def __init__(self, easyeda_cp_cad_data, lcsc_id, download_raw_3d_model=True):
        self.input = easyeda_cp_cad_data
        self.download_raw_3d_model = download_raw_3d_model
        self.output = self.create_3d_model(lcsc_id)

    def create_3d_model(self, lcsc_id):
        ee_data = self.input["packageDetail"]["dataStr"]["shape"] if isinstance(self.input, dict) else self.input

        model_3d_info = self.get_3d_model_info(ee_data=ee_data)
        if model_3d_info:
            model_3d = self.parse_3d_model_info(info=model_3d_info)
            if self.download_raw_3d_model:
                model_3d.raw_obj, model_3d.step = EasyedaApi().get_raw_3d_model_obj(uuid=model_3d.uuid)
            return model_3d

        logger.warning(W_EEDA3D+"No 3D model available for `{}`".format(lcsc_id))
        return None

    def get_3d_model_info(self, ee_data: str):
        for line in ee_data:
            ee_designator = line.split("~")[0]
            if ee_designator == "SVGNODE":
                raw_json = line.split("~")[1:][0]
                return json.loads(raw_json)["attrs"]
        return {}

    def parse_3d_model_info(self, info: dict):
        return Ee3dModel(
            name=info["title"],
            uuid=info["uuid"],
            translation=Ee3dModelBase(
                x=info["c_origin"].split(",")[0],
                y=info["c_origin"].split(",")[1],
                z=info["z"],
            ),
            rotation=Ee3dModelBase(*(info["c_rotation"].split(","))),
        )


def get_materials(obj_data: str):

    material_regex = "newmtl .*?endmtl"
    matches = re.findall(pattern=material_regex, string=obj_data, flags=re.DOTALL)

    materials = {}
    for c, match in enumerate(matches):
        material = {"ambient_intensity": 0.2, "shininess": 0.5, "transparency": 0}
        material["name"] = "MATERIAL_"+str(c+1)
        for value in match.splitlines():
            if value.startswith("newmtl"):
                material_id = value.split(" ")[1]
            elif value.startswith("Ka"):
                material["ambient_color"] = value.split(" ")[1:]
            elif value.startswith("Kd"):
                material["diffuse_color"] = value.split(" ")[1:]
            elif value.startswith("Ks"):
                material["specular_color"] = value.split(" ")[1:]
            # elif value.startswith("d"):
            #   material["transparency"] = value.split(" ")[1]

        materials[material_id] = material
    return materials


def get_vertices(obj_data: str):
    vertices_regex = "v (.*?)\n"
    matches = re.findall(pattern=vertices_regex, string=obj_data, flags=re.DOTALL)

    return [" ".join([str(round(float(coord) / 2.54, 4)) for coord in vertice.split(" ")]) for vertice in matches]


def map_materials(mats):
    """ An heuristic to map generic colors to the ones used by KiCad """
    # Look for grey, black and yellow materials.
    greys = []
    blacks = []
    golden = []
    for id, mat in mats.items():
        r, g, b = map(float, mat['diffuse_color'])
        if abs(r-g) < 0.02 and abs(g-b) < 0.02:
            # Same RGB components
            if r < 0.97 and r > 0.6:
                # In the 0.6 - 0.97 range
                greys.append(id)
            if r < 0.3 and r > 0.1:
                blacks.append(id)
        elif r > 0.8 and g < 0.8 and g > 0.7 and b < 0.5:
            golden.append(id)
        if (r == 0 and g == 0 and b == 0) or (r == 1 and g == 1 and b == 1):
            mat['transparency'] = 1
        else:
            mat['transparency'] = 0
        mat['diffuse_color'] = ('%5.3f' % r, '%5.3f' % g, '%5.3f' % b)
        mat['specular_color'] = tuple('%5.3f' % float(c) for c in mat['specular_color'])
    # Use greys for the pins and metal body
    c_greys = len(greys)
    if c_greys == 1:
        mats[greys[0]] = MATERIAL_PIN_01
    elif c_greys > 1:
        # More than one grey, sort by specular level
        greys = sorted(greys, key=lambda x: mats[x]['specular_color'][0], reverse=True)
        mats[greys[0]] = MATERIAL_PIN_01
        mats[greys[1]] = MATERIAL_MET_01
    # Use black for the plastic body
    c_blacks = len(blacks)
    if c_blacks == 1:
        mats[blacks[0]] = MATERIAL_EPOXY_04
    elif c_blacks > 1:
        blacks = sorted(blacks, key=lambda x: mats[x]['diffuse_color'][0], reverse=True)
        mats[blacks[0]] = MATERIAL_EPOXY_04
        # for c, b in enumerate(blacks):
        #     mat_map[b] = 'IC-BODY-EPOXY-%02d' % c+1
    # Use yellow for golden pins
    if golden:
        mats[golden[0]] = MATERIAL_PIN_02


def generate_wrl_model(model_3d: Ee3dModel):
    materials = get_materials(obj_data=model_3d.raw_obj)
    vertices = get_vertices(obj_data=model_3d.raw_obj)
    map_materials(materials)

    raw_wrl = VRML_HEADER
    # Define all the materials
    for mat in materials.values():
        mat_str = textwrap.dedent(
            f"""
            Shape {{
                appearance Appearance {{
                    material DEF {mat['name']} Material {{
                        ambientIntensity {mat['ambient_intensity']}
                        diffuseColor {' '.join(mat['diffuse_color'])}
                        specularColor {' '.join(mat['specular_color'])}
                        shininess {mat['shininess']}
                        transparency {mat['transparency']}
                    }}
                }}
            }}"""
        )
        raw_wrl += mat_str
    # Define the shapes
    shapes = model_3d.raw_obj.split("usemtl")[1:]
    for shape in shapes:
        lines = shape.splitlines()
        material_id = materials[lines[0].replace(" ", "")]["name"]
        index_counter = 0
        link_dict = {}
        coord_index = []
        points = []
        for line in lines[1:]:
            if len(line) > 0:
                face = [int(index) for index in line.replace("//", "").split(" ")[1:]]
                face_index = []
                for index in face:
                    if index not in link_dict:
                        link_dict[index] = index_counter
                        face_index.append(str(index_counter))
                        points.append(vertices[index - 1])
                        index_counter += 1
                    else:
                        face_index.append(str(link_dict[index]))
                face_index.append("-1")
                coord_index.append(",".join(face_index) + ",")
        points.insert(-1, points[-1])

        shape_str = textwrap.dedent(
            f"""
            Shape {{
                geometry IndexedFaceSet {{
                    ccw TRUE
                    solid FALSE
                    coord DEF co Coordinate {{
                        point [
                            {(", ").join(points)}
                        ]
                    }}
                    coordIndex [
                        {"".join(coord_index)}
                    ]
                }}
                appearance Appearance {{material USE {material_id}}}
            }}"""
        )

        raw_wrl += shape_str

    return Ki3dModel(translation=None, rotation=None, name=model_3d.name, raw_wrl=raw_wrl)


class Exporter3dModelKicad:
    def __init__(self, model_3d):
        self.input = model_3d
        self.output = generate_wrl_model(model_3d=model_3d) if model_3d and model_3d.raw_obj else None
        self.output_step = model_3d.step if model_3d and model_3d.step else None

    def export(self, lib_path, fname=None):
        name_wrl = name_step = None
        # Export the WRL
        if self.output is not None:
            name_wrl = fname
            if name_wrl is None:
                name_wrl = self.output.name+'.wrl'
            name_wrl = os.path.join(lib_path, name_wrl)
            with open(name_wrl, "w", encoding="utf-8") as my_lib:
                my_lib.write(self.output.raw_wrl)
        # Export the STEP
        if self.output_step is not None:
            name_step = fname
            if name_step is None:
                name_step = (self.output.name if self.output else self.input.uuid)+'.step'
            else:
                name_step = os.path.splitext(name_step)[0]+'.step'
            name_step = os.path.join(lib_path, name_step)
            with open(name_step, "wb") as my_lib:
                my_lib.write(self.output_step)
        return name_wrl or name_step


def download_easyeda_3d_model(lcsc_id, dest_path, fname=None):
    api = EasyedaApi()
    data = api.get_cad_data_of_component(lcsc_id)
    if not data:
        return None
    importer = Easyeda3dModelImporter(data, lcsc_id)
    exporter = Exporter3dModelKicad(model_3d=importer.output)
    os.makedirs(dest_path, exist_ok=True)
    return exporter.export(dest_path, fname)
