use crate::Shape;
use crate::parser::mmcif::Chain;
use crate::parser::mmcif::MmCif;
use crate::parser::mmcif::Residue;
use crate::parser::mmcif::ResidueType;
use crate::parser::mmcif::SecondaryStructure;
use crate::utils::{MeshData, VisualShape, VisualStyle};
use glam::Vec3;
use na_seq::AtomTypeInRes;
use serde::{Deserialize, Serialize};

#[derive(Serialize, Deserialize, Debug, Clone)]
pub struct Protein {
    pub chains: Vec<Chain>,
    pub center: Vec3,
}

impl Protein {
    pub fn new(mmcif: MmCif) -> Self {
        let mut chains = Vec::new();
        let mut centers = Vec::new();

        for chain in mmcif.chains {
            let mut residues = Vec::new();
            for residue_id in chain.residue_sns {
                let residue = &mmcif.residues[centers.len()];
                let res_type: ResidueType = residue.res_type.clone();
                let mut ca_opt = None;
                let mut cb_opt = None;
                for atom_id in &residue.atom_sns {
                    let atom = &mmcif.atoms[*atom_id as usize - 1];
                    if let Some(atom_type_in_res) = &atom.type_in_res {
                        if *atom_type_in_res == AtomTypeInRes::CA {
                            ca_opt = Some(atom.posit);
                        }
                        if *atom_type_in_res == AtomTypeInRes::CB {
                            cb_opt = Some(atom.posit);
                        }
                    }
                }

                if ca_opt.is_none() {
                    println!(
                        "No CA atom found for chain {} residue {}",
                        chain.id, residue_id
                    );
                    continue;
                }

                let ca = ca_opt.expect("No CA atom found!");

                centers.push(Vec3 {
                    x: ca.x as f32,
                    y: ca.y as f32,
                    z: ca.z as f32,
                });

                residues.push(Residue {
                    residue_type: res_type,
                    ca: Vec3 {
                        x: ca_opt.unwrap().x as f32
                            + match chain.id.as_str() {
                                "A" => 0.0,
                                "B" => 1.0,
                                _ => 2.0,
                            },
                        y: ca_opt.unwrap().y as f32,
                        z: ca_opt.unwrap().z as f32,
                    },
                    cb: match cb_opt {
                        Some(pos) => Some(Vec3 {
                            x: pos.x as f32,
                            y: pos.y as f32,
                            z: pos.z as f32,
                        }),
                        None => None,
                    },
                    index: residue_id as usize,
                    ss: SecondaryStructure::Unknown,
                });
            }
            chains.push(Chain {
                residues: residues,
                id: chain.id,
            });
        }

        let mut center = Vec3::ZERO;
        for c in &centers {
            center += c;
        }
        center = center / (centers.len() as f32);

        Protein {
            chains: chains,
            center: center,
        }
    }
}

impl VisualShape for Protein {
    fn style_mut(&mut self) -> &mut VisualStyle {
        unimplemented!()
    }
}

impl Protein {
    pub fn get_center(&self) -> [f32; 3] {
        [self.center.x, self.center.y, self.center.z]
    }

    pub fn centered(mut self) -> Self {
        let center = Vec3 {
            x: self.center.x,
            y: self.center.y,
            z: self.center.z,
        };
        for chain in &mut self.chains {
            for residue in &mut chain.residues {
                residue.ca -= center;
                if let Some(cb) = residue.cb {
                    residue.cb = Some(cb - center);
                }
            }
        }
        self.center = Vec3::ZERO;
        self
    }

    fn catmull_rom_chain(&self, positions: &[Vec3]) -> Vec<Vec3> {
        if positions.len() < 2 {
            return positions.to_vec();
        }

        let n = positions.len();
        let mut path = Vec::with_capacity(n * 5 + 1); // 每残基 5 个点

        // ChimeraX 默认参数：tension = 0.5（其实就是标准 Catmull-Rom）
        const SEGMENTS_PER_RESIDUE: usize = 5;

        for i in 0..n - 1 {
            let p0 = if i > 0 {
                positions[i - 1]
            } else {
                positions[i]
            };
            let p1 = positions[i];
            let p2 = positions[i + 1];
            let p3 = if i + 2 < n {
                positions[i + 2]
            } else {
                positions[i + 1]
            };

            for j in 0..=SEGMENTS_PER_RESIDUE {
                let t = j as f32 / SEGMENTS_PER_RESIDUE as f32;
                let pos = catmull_rom(p0, p1, p2, p3, t);
                path.push(pos);
            }
        }

        // 最后一个残基的终点（避免重复）
        path.push(positions[n - 1]);
        path
    }

    pub fn to_mesh(&self, scale: f32) -> MeshData {
        let mut final_mesh = MeshData::default();
        for chain in &self.chains {
            let mut mesh = MeshData::default();

            let residues: Vec<&Residue> = chain
                .residues
                .iter()
                .filter(|r| r.ca.length_squared() > 1e-6)
                .collect();

            let ca_positions: Vec<Vec3> = residues.iter().map(|r| r.ca).collect();

            if ca_positions.len() < 2 {
                println!("chain {} has less than 2 residues, skipping", chain.id);
                continue;
            }

            // 直接生成平滑路径 + 切线（不再依赖 splines）
            let path = self.catmull_rom_chain(&ca_positions);

            let mut centers = Vec::with_capacity(path.len());
            let mut tangents = Vec::with_capacity(path.len());
            let mut normals = Vec::with_capacity(path.len());

            let n = path.len();
            for i in 0..n {
                let pos = path[i];
                centers.push(pos);

                // 计算切线（中心差分 + 边界特殊处理）
                let tan = if i == 0 {
                    (path[1] - path[0]).normalize()
                } else if i == n - 1 {
                    (path[n - 1] - path[n - 2]).normalize()
                } else {
                    (path[i + 1] - path[i - 1]).normalize()
                };
                tangents.push(tan);

                // 初始化 normal
                let mut n = if tan.dot(Vec3::Z).abs() > 0.98 {
                    Vec3::Y
                } else {
                    Vec3::Z
                };
                n = (n - n.dot(tan) * tan).normalize();
                if n.length_squared() < 1e-6 {
                    n = tan.any_orthonormal_vector();
                }
                normals.push(n);
            }

            // Parallel transport（最小扭转）
            for i in 1..n {
                let t = tangents[i];
                let prev_n = normals[i - 1];
                let b = t.cross(prev_n).normalize_or_zero();
                let new_n = b.cross(t).normalize();

                if new_n.dot(prev_n) < 0.0 {
                    normals[i] = -new_n;
                } else {
                    normals[i] = new_n;
                }
            }

            // sections 和之前一样

            let sections: Vec<RibbonXSection> = residues
                .iter()
                .map(|r| match r.ss {
                    SecondaryStructure::Helix => helix_section(),
                    SecondaryStructure::Sheet => sheet_section(),
                    _ => coil_section(),
                })
                .collect();
            // let sections: Vec<RibbonXSection> = residues.iter().map(|r| helix_section()).collect();

            // extrusion（保持你最新的 extrude_ribbon_corrected）
            self.extrude_ribbon_corrected(&centers, &tangents, &normals, &sections, 5, &mut mesh);

            // 缩放 + 颜色
            for v in &mut mesh.vertices {
                *v = (Vec3::from_array(*v) * scale).to_array();
            }
            mesh.colors = Some(vec![[1.0, 1.0, 1.0, 1.0]; mesh.vertices.len()]);

            final_mesh.append(&mesh);
        }
        final_mesh
    }

    // 完全修正版的 extrusion（不再有任何越界、箭头方向、端盖问题）
    fn extrude_ribbon_corrected(
        &self,
        centers: &[Vec3],
        tangents: &[Vec3],
        normals: &[Vec3],
        sections: &[RibbonXSection],
        pts_per_res: usize,
        mesh: &mut MeshData,
    ) {
        let base_v = mesh.vertices.len() as u32;

        for (seg, xs) in sections.iter().enumerate() {
            let start = seg * pts_per_res;
            // 注意：最后一个 segment 多取 pts_per_res 个点（共享下一个 segment 的起点）
            let end = if seg + 1 < sections.len() {
                (seg + 1) * pts_per_res + 1
            } else {
                centers.len() // 最后一个正好到结尾
            };

            let cap_front = seg == 0;
            let cap_back = seg + 1 == sections.len();

            let (coords, arrow_back) = if xs.arrow_coords.is_some() {
                (&xs.coords[..], xs.arrow_coords.as_deref())
            } else {
                (&xs.coords[..], None)
            };

            self.extrude_one_segment(
                &centers[start..end],
                &tangents[start..end],
                &normals[start..end],
                coords,
                arrow_back,
                cap_front,
                cap_back,
                mesh,
            );
        }

        // 统一偏移索引
        for idx in &mut mesh.indices {
            *idx += base_v;
        }
    }

    #[allow(clippy::too_many_arguments)]
    fn extrude_one_segment(
        &self,
        centers: &[Vec3],
        tangents: &[Vec3],
        normals: &[Vec3],
        coords: &[[f32; 2]],
        arrow_back: Option<&[[f32; 2]]>,
        cap_front: bool,
        cap_back: bool,
        mesh: &mut MeshData,
    ) {
        let n_ring = coords.len();
        let n_pts = centers.len();
        let base = mesh.vertices.len() as u32;

        // binormal = T × N
        let binormals: Vec<Vec3> = tangents
            .iter()
            .zip(normals)
            .map(|(t, n)| t.cross(*n).normalize_or_zero())
            .collect();

        // 生成顶点
        for i in 0..n_pts {
            let c = centers[i];
            let n = normals[i];
            let b = binormals[i];

            // Sheet 箭头只在前半段渐变（ChimeraX 就是这么做的）
            let arrow_t = if arrow_back.is_some() {
                let local_t = i as f32 / (n_pts.saturating_sub(1)) as f32;
                (local_t * 2.0).min(1.0) // 0~0.5 → 1.0~0.0
            } else {
                0.0
            };
            let arrow_factor = if arrow_t > 0.0 { 1.0 - arrow_t } else { 1.0 };

            for ring_idx in 0..n_ring {
                let mut off = coords[ring_idx];

                if arrow_factor < 1.0 && arrow_back.is_some() {
                    let back = arrow_back.unwrap()[ring_idx];
                    off[0] = off[0] + (back[0] - off[0]) * (1.0 - arrow_factor);
                    off[1] = off[1] + (back[1] - off[1]) * (1.0 - arrow_factor);
                }

                let pos = c + n * off[0] + b * off[1];
                let nor = (n * off[0] + b * off[1]).normalize();

                mesh.vertices.push(pos.to_array());
                mesh.normals.push(nor.to_array());
            }
        }

        // 四边形 → 三角形
        for i in 0..n_pts - 1 {
            for r in 0..n_ring {
                let r_next = (r + 1) % n_ring;

                let i0 = i * n_ring + r;
                let i1 = i * n_ring + r_next;
                let j0 = (i + 1) * n_ring + r;
                let j1 = (i + 1) * n_ring + r_next;

                let a = base + i0 as u32;
                let b = base + i1 as u32;
                let c = base + j1 as u32;
                let d = base + j0 as u32;

                mesh.indices.extend_from_slice(&[a, b, d]);
                // mesh.indices.extend_from_slice(&[a, d, b]);
                mesh.indices.extend_from_slice(&[b, c, d]);
            }
        }

        // 端盖（扇形三角化）
        let mut cap = |start_pt: usize, tangent: Vec3, outward: bool| {
            let center = base + (start_pt * n_ring) as u32;
            let normal = if outward { tangent } else { -tangent };

            // 简单星形三角化（足够光滑）
            for r in 1..n_ring - 1 {
                let a = center;
                let b = center + r as u32;
                let c = center + (r + 1) as u32;

                let v_ab: Vec3 = Vec3::from_array(mesh.vertices[b as usize])
                    - Vec3::from_array(mesh.vertices[a as usize]);
                let v_ac: Vec3 = Vec3::from_array(mesh.vertices[c as usize])
                    - Vec3::from_array(mesh.vertices[a as usize]);

                if normal.dot(v_ab.cross(v_ac)) > 0.0 {
                    mesh.indices.extend_from_slice(&[a, c, b]);
                } else {
                    mesh.indices.extend_from_slice(&[a, b, c]);
                }
            }
        };

        if cap_front {
            cap(0, tangents[0], false);
        }
        if cap_back {
            cap(n_pts - 1, tangents[n_pts - 1], true);
        }
    }
}

// 标准 Catmull-Rom 公式（ChimeraX、Mol*、PyMOL、VMD 全都用这个）
#[inline(always)]
fn catmull_rom(p0: Vec3, p1: Vec3, p2: Vec3, p3: Vec3, t: f32) -> Vec3 {
    let t2 = t * t;
    let t3 = t2 * t;

    // Catmull-Rom 系数（tension = 0.5）
    let c0 = -0.5 * t3 + t2 - 0.5 * t;
    let c1 = 1.5 * t3 - 2.5 * t2 + 1.0;
    let c2 = -1.5 * t3 + 2.0 * t2 + 0.5 * t;
    let c3 = 0.5 * t3 - 0.5 * t2;

    p0 * c0 + p1 * c1 + p2 * c2 + p3 * c3
}

impl Into<Shape> for Protein {
    fn into(self) -> Shape {
        Shape::Protein(self)
    }
}

// ---------- 下面是 ChimeraX 完全一致的 cross-section 定义 ----------
fn helix_section() -> RibbonXSection {
    // 圆形（实际上是 32 边形近似）+ 箭头（Sheet 用）
    let mut coords = Vec::new();
    let n = 32;
    for i in 0..n {
        let a = (i as f32) / (n as f32) * std::f32::consts::TAU;
        coords.push([a.cos(), a.sin()]);
    }
    RibbonXSection::smooth_circle(&coords)
}

fn sheet_section() -> RibbonXSection {
    // ChimeraX 默认箭头：前半段宽 1.0 → 0.0，后半段保持 0.0
    let base_arrow = helix_section().arrow(1.0, 1.0, 0.0, 1.2); // 经典箭头比例
    base_arrow
}

fn coil_section() -> RibbonXSection {
    // Coil 用细一点的圆管
    helix_section().scale(0.4, 0.4)
}

struct RibbonXSection {
    coords: Vec<[f32; 2]>,               // 基础 2D 轮廓
    arrow_coords: Option<Vec<[f32; 2]>>, // 为 Sheet 箭头准备的第二套轮廓
    _smooth: bool,
}

impl RibbonXSection {
    fn smooth_circle(coords: &[[f32; 2]]) -> Self {
        Self {
            coords: coords.to_vec(),
            arrow_coords: None,
            _smooth: true,
        }
    }

    fn scale(mut self, sx: f32, sy: f32) -> Self {
        for c in &mut self.coords {
            c[0] *= sx;
            c[1] *= sy;
        }
        if let Some(ac) = self.arrow_coords.as_mut() {
            for c in ac {
                c[0] *= sx;
                c[1] *= sy;
            }
        }
        self
    }

    fn arrow(mut self, sx1: f32, sy1: f32, sx2: f32, sy2: f32) -> Self {
        let mut back = self.coords.clone();
        for c in &mut self.coords {
            c[0] *= sx1;
            c[1] *= sy1;
        }
        for c in &mut back {
            c[0] *= sx2;
            c[1] *= sy2;
        }
        self.arrow_coords = Some(back);
        self
    }
}
