#![allow(async_fn_in_trait)]

use serde::Deserialize;
use serde::Serialize;

use crate::ExecutorContext;
use crate::KclErrorWithOutputs;
use crate::front::Plane;
use crate::frontend::api::Expr;
use crate::frontend::api::FileId;
use crate::frontend::api::Number;
use crate::frontend::api::ObjectId;
use crate::frontend::api::ProjectId;
use crate::frontend::api::SceneGraph;
use crate::frontend::api::SceneGraphDelta;
use crate::frontend::api::SourceDelta;
use crate::frontend::api::Version;

pub type ExecResult<T> = std::result::Result<T, KclErrorWithOutputs>;

/// Information about a newly created segment for batch operations
#[derive(Debug, Clone)]
pub struct NewSegmentInfo {
    pub segment_id: ObjectId,
    pub start_point_id: ObjectId,
    pub end_point_id: ObjectId,
    pub center_point_id: Option<ObjectId>,
}

pub trait SketchApi {
    /// Execute the sketch in mock mode, without changing anything. This is
    /// useful after editing segments, and the user releases the mouse button.
    async fn execute_mock(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn new_sketch(
        &mut self,
        ctx: &ExecutorContext,
        project: ProjectId,
        file: FileId,
        version: Version,
        args: SketchCtor,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta, ObjectId)>;

    // Enters sketch mode
    async fn edit_sketch(
        &mut self,
        ctx: &ExecutorContext,
        project: ProjectId,
        file: FileId,
        version: Version,
        sketch: ObjectId,
    ) -> ExecResult<SceneGraphDelta>;

    async fn exit_sketch(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
    ) -> ExecResult<SceneGraph>;

    async fn delete_sketch(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn add_segment(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        segment: SegmentCtor,
        label: Option<String>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn edit_segments(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        segments: Vec<ExistingSegmentCtor>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn delete_objects(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        constraint_ids: Vec<ObjectId>,
        segment_ids: Vec<ObjectId>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn add_constraint(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        constraint: Constraint,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn chain_segment(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        previous_segment_end_point_id: ObjectId,
        segment: SegmentCtor,
        label: Option<String>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn edit_constraint(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        constraint_id: ObjectId,
        value_expression: String,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    async fn edit_distance_constraint_label_position(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        constraint_id: ObjectId,
        label_position: Point2d<Number>,
        anchor_segment_ids: Vec<ObjectId>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    /// Batch operations for split segment: edit segments, add constraints, delete objects.
    /// All operations are applied to a single AST and execute_after_edit is called once at the end.
    /// new_segment_info contains the IDs from the segment(s) added in a previous step.
    #[allow(clippy::too_many_arguments)]
    async fn batch_split_segment_operations(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        edit_segments: Vec<ExistingSegmentCtor>,
        add_constraints: Vec<Constraint>,
        delete_constraint_ids: Vec<ObjectId>,
        new_segment_info: NewSegmentInfo,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;

    /// Batch operations for tail-cut trim: edit a segment, add coincident constraints,
    /// delete constraints, and execute once.
    #[allow(clippy::too_many_arguments)]
    async fn batch_tail_cut_operations(
        &mut self,
        ctx: &ExecutorContext,
        version: Version,
        sketch: ObjectId,
        edit_segments: Vec<ExistingSegmentCtor>,
        add_constraints: Vec<Constraint>,
        delete_constraint_ids: Vec<ObjectId>,
        additional_edited_segment_ids: Vec<ObjectId>,
    ) -> ExecResult<(SourceDelta, SceneGraphDelta)>;
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiSketch")]
pub struct Sketch {
    pub args: SketchCtor,
    pub plane: ObjectId,
    pub segments: Vec<ObjectId>,
    pub constraints: Vec<ObjectId>,
}

/// Arguments for creating a new sketch. This is similar to the constructor of
/// other kinds of objects in that it is the inputs to the sketch, not the
/// outputs.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct SketchCtor {
    /// The sketch surface.
    pub on: Plane,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiPoint")]
pub struct Point {
    pub position: Point2d<Number>,
    pub ctor: Option<PointCtor>,
    pub owner: Option<ObjectId>,
    pub freedom: Freedom,
    pub constraints: Vec<ObjectId>,
}

impl Point {
    /// The freedom of this point.
    pub fn freedom(&self) -> Freedom {
        self.freedom
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub enum Freedom {
    Free,
    Fixed,
    Conflict,
}

impl Freedom {
    /// Merges two Freedom values. For example, a point has a solver variable
    /// for each dimension, x and y. If one dimension is `Free` and the other is
    /// `Fixed`, the point overall is `Free` since it isn't fully constrained.
    /// `Conflict` infects the most, followed by `Free`. An object must be fully
    /// `Fixed` to be `Fixed` overall.
    pub fn merge(self, other: Self) -> Self {
        match (self, other) {
            (Self::Conflict, _) | (_, Self::Conflict) => Self::Conflict,
            (Self::Free, _) | (_, Self::Free) => Self::Free,
            (Self::Fixed, Self::Fixed) => Self::Fixed,
        }
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiSegment")]
#[serde(tag = "type")]
pub enum Segment {
    Point(Point),
    Line(Line),
    Arc(Arc),
    Circle(Circle),
    ControlPointSpline(ControlPointSpline),
}

impl Segment {
    /// What kind of geometry is this (point, line, arc, etc)
    /// Suitable for use in user-facing messages.
    pub fn human_friendly_kind_with_article(&self) -> &'static str {
        match self {
            Self::Point(_) => "a Point",
            Self::Line(_) => "a Line",
            Self::Arc(_) => "an Arc",
            Self::Circle(_) => "a Circle",
            Self::ControlPointSpline(_) => "a Control Point Spline",
        }
    }

    /// Compute the overall freedom of this segment. For geometry types (Line,
    /// Arc, Circle) this looks up and merges the freedom of their constituent
    /// points. For points, returns the point's own freedom directly.
    /// Returns `None` if a required point lookup failed.
    pub fn freedom(&self, lookup: impl Fn(ObjectId) -> Option<Freedom>) -> Option<Freedom> {
        match self {
            Self::Point(p) => Some(p.freedom()),
            Self::Line(l) => l.freedom(&lookup),
            Self::Arc(a) => a.freedom(&lookup),
            Self::Circle(c) => c.freedom(&lookup),
            Self::ControlPointSpline(s) => s.freedom(&lookup),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct ExistingSegmentCtor {
    pub id: ObjectId,
    pub ctor: SegmentCtor,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
#[serde(tag = "type")]
pub enum SegmentCtor {
    Point(PointCtor),
    Line(LineCtor),
    Arc(ArcCtor),
    Circle(CircleCtor),
    ControlPointSpline(ControlPointSplineCtor),
}

impl SegmentCtor {
    /// What kind of geometry is this (point, line, arc, etc)
    /// Suitable for use in user-facing messages.
    pub fn human_friendly_kind_with_article(&self) -> &'static str {
        match self {
            Self::Point(_) => "a Point constructor",
            Self::Line(_) => "a Line constructor",
            Self::Arc(_) => "an Arc constructor",
            Self::Circle(_) => "a Circle constructor",
            Self::ControlPointSpline(_) => "a Control Point Spline constructor",
        }
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct PointCtor {
    pub position: Point2d<Expr>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiPoint2d")]
pub struct Point2d<U: std::fmt::Debug + Clone + ts_rs::TS> {
    pub x: U,
    pub y: U,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiLine")]
pub struct Line {
    pub start: ObjectId,
    pub end: ObjectId,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub owner: Option<ObjectId>,
    // Invariant: Line or MidPointLine
    pub ctor: SegmentCtor,
    // The constructor is applicable if changing the values of the constructor will change the rendering
    // of the segment (modulo multiple valid solutions). I.e., whether the object is constrained with
    // respect to the constructor inputs.
    // The frontend should only display handles for the constructor inputs if the ctor is applicable.
    // (Or because they are the (locked) start/end of the segment).
    pub ctor_applicable: bool,
    pub construction: bool,
}

impl Line {
    /// Compute the overall freedom of this line by merging the freedom of its
    /// start and end points. Returns `None` if a point lookup failed.
    pub fn freedom(&self, lookup: impl Fn(ObjectId) -> Option<Freedom>) -> Option<Freedom> {
        let start = lookup(self.start)?;
        let end = lookup(self.end)?;
        Some(start.merge(end))
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct LineCtor {
    pub start: Point2d<Expr>,
    pub end: Point2d<Expr>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub construction: Option<bool>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiStartOrEnd")]
#[serde(tag = "type")]
pub enum StartOrEnd<T> {
    Start(T),
    End(T),
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiArc")]
pub struct Arc {
    pub start: ObjectId,
    pub end: ObjectId,
    pub center: ObjectId,
    // Invariant: Arc
    pub ctor: SegmentCtor,
    pub ctor_applicable: bool,
    pub construction: bool,
}

impl Arc {
    /// Compute the overall freedom of this arc by merging the freedom of its
    /// start, end, and center points. Returns `None` if a point lookup failed.
    pub fn freedom(&self, lookup: impl Fn(ObjectId) -> Option<Freedom>) -> Option<Freedom> {
        let start = lookup(self.start)?;
        let end = lookup(self.end)?;
        let center = lookup(self.center)?;
        Some(start.merge(end).merge(center))
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct ArcCtor {
    pub start: Point2d<Expr>,
    pub end: Point2d<Expr>,
    pub center: Point2d<Expr>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub construction: Option<bool>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiCircle")]
pub struct Circle {
    pub start: ObjectId,
    pub center: ObjectId,
    // Invariant: Circle
    pub ctor: SegmentCtor,
    pub ctor_applicable: bool,
    pub construction: bool,
}

impl Circle {
    /// Compute the overall freedom of this circle by merging the freedom of its
    /// start and center points. Returns `None` if a point lookup failed.
    pub fn freedom(&self, lookup: impl Fn(ObjectId) -> Option<Freedom>) -> Option<Freedom> {
        let start = lookup(self.start)?;
        let center = lookup(self.center)?;
        Some(start.merge(center))
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct CircleCtor {
    pub start: Point2d<Expr>,
    pub center: Point2d<Expr>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub construction: Option<bool>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiControlPointSpline")]
pub struct ControlPointSpline {
    pub controls: Vec<ObjectId>,
    pub degree: u32,
    pub ctor: SegmentCtor,
    pub ctor_applicable: bool,
    pub construction: bool,
}

impl ControlPointSpline {
    /// Compute the overall freedom of this spline by merging the freedom of its
    /// control points. Returns `None` if any required point lookup failed.
    pub fn freedom(&self, lookup: impl Fn(ObjectId) -> Option<Freedom>) -> Option<Freedom> {
        let mut controls = self.controls.iter();
        let first = lookup(*controls.next()?)?;
        let merged = controls.try_fold(first, |acc, id| lookup(*id).map(|freedom| acc.merge(freedom)))?;
        Some(merged)
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct ControlPointSplineCtor {
    pub points: Vec<Point2d<Expr>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    #[ts(optional)]
    pub construction: Option<bool>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", rename = "ApiConstraint")]
#[serde(tag = "type")]
// When adding a new constraint type, check trim compatibility. New constraints
// can break trim in unexpected ways, especially when endpoints are edited,
// segments are split, or constraints are migrated. Try the trim tool on sketches
// using the new constraint, and talk to Kurt, Max, or a mechanical engineer if
// the intended trim behavior is unclear.
pub enum Constraint {
    Coincident(Coincident),
    Distance(Distance),
    Angle(Angle),
    Diameter(Diameter),
    EqualRadius(EqualRadius),
    Fixed(Fixed),
    HorizontalDistance(Distance),
    VerticalDistance(Distance),
    Horizontal(Horizontal),
    LinesEqualLength(LinesEqualLength),
    Midpoint(Midpoint),
    Parallel(Parallel),
    Perpendicular(Perpendicular),
    Radius(Radius),
    Symmetric(Symmetric),
    Tangent(Tangent),
    Vertical(Vertical),
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Coincident {
    pub segments: Vec<ConstraintSegment>,
}

impl Coincident {
    pub fn get_segments(&self) -> Vec<ObjectId> {
        self.segments
            .iter()
            .filter_map(|segment| match segment {
                ConstraintSegment::Segment(id) => Some(*id),
                ConstraintSegment::Origin(_) => None,
            })
            .collect()
    }

    pub fn segment_ids(&self) -> impl Iterator<Item = ObjectId> + '_ {
        self.segments.iter().filter_map(|segment| match segment {
            ConstraintSegment::Segment(id) => Some(*id),
            ConstraintSegment::Origin(_) => None,
        })
    }

    pub fn contains_segment(&self, segment_id: ObjectId) -> bool {
        self.segment_ids().any(|id| id == segment_id)
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
#[serde(untagged)]
pub enum ConstraintSegment {
    Segment(ObjectId),
    Origin(OriginLiteral),
}

impl ConstraintSegment {
    pub const ORIGIN: Self = Self::Origin(OriginLiteral::Origin);
}

#[derive(Debug, Clone, Copy, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
#[serde(rename_all = "SCREAMING_SNAKE_CASE")]
pub enum OriginLiteral {
    Origin,
}

impl From<ObjectId> for ConstraintSegment {
    fn from(value: ObjectId) -> Self {
        Self::Segment(value)
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Distance {
    pub points: Vec<ConstraintSegment>,
    pub distance: Number,
    #[serde(rename = "labelPosition")]
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[ts(rename = "labelPosition")]
    #[ts(optional)]
    pub label_position: Option<Point2d<Number>>,
    pub source: ConstraintSource,
}

impl Distance {
    pub fn point_ids(&self) -> impl Iterator<Item = ObjectId> + '_ {
        self.points.iter().filter_map(|point| match point {
            ConstraintSegment::Segment(id) => Some(*id),
            ConstraintSegment::Origin(_) => None,
        })
    }

    pub fn contains_point(&self, point_id: ObjectId) -> bool {
        self.point_ids().any(|id| id == point_id)
    }
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Angle {
    pub lines: Vec<ObjectId>,
    pub angle: Number,
    pub source: ConstraintSource,
}

#[derive(Debug, Clone, Default, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct ConstraintSource {
    pub expr: String,
    pub is_literal: bool,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Radius {
    pub arc: ObjectId,
    pub radius: Number,
    #[serde(rename = "labelPosition")]
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[ts(rename = "labelPosition")]
    #[ts(optional)]
    pub label_position: Option<Point2d<Number>>,
    #[serde(default)]
    pub source: ConstraintSource,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Diameter {
    pub arc: ObjectId,
    pub diameter: Number,
    #[serde(rename = "labelPosition")]
    #[serde(default, skip_serializing_if = "Option::is_none")]
    #[ts(rename = "labelPosition")]
    #[ts(optional)]
    pub label_position: Option<Point2d<Number>>,
    #[serde(default)]
    pub source: ConstraintSource,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", optional_fields)]
pub struct EqualRadius {
    pub input: Vec<ObjectId>,
}

/// Multiple fixed constraints, allowing callers to add fixed constraints on
/// multiple points at once.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Fixed {
    pub points: Vec<FixedPoint>,
}

/// A fixed constraint on a single point.
#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct FixedPoint {
    pub point: ObjectId,
    pub position: Point2d<Number>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
#[serde(untagged)]
pub enum Horizontal {
    Line { line: ObjectId },
    Points { points: Vec<ConstraintSegment> },
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct LinesEqualLength {
    pub lines: Vec<ObjectId>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
pub struct Midpoint {
    pub point: ConstraintSegment,
    #[serde(alias = "line")]
    pub segment: ObjectId,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts")]
#[serde(untagged)]
pub enum Vertical {
    Line { line: ObjectId },
    Points { points: Vec<ConstraintSegment> },
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", optional_fields)]
pub struct Parallel {
    pub lines: Vec<ObjectId>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", optional_fields)]
pub struct Perpendicular {
    pub lines: Vec<ObjectId>,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", optional_fields)]
pub struct Symmetric {
    pub input: Vec<ObjectId>,
    pub axis: ObjectId,
}

#[derive(Debug, Clone, PartialEq, Deserialize, Serialize, ts_rs::TS)]
#[ts(export, export_to = "FrontendApi.ts", optional_fields)]
pub struct Tangent {
    pub input: Vec<ObjectId>,
}
