// Copyright 2025 SkyReels SDK Contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use crate::error::{Error, Result};
use crate::client::SkyReelsClient;
use crate::text2video::{Mode, TaskStatus, SubmitResponse, TaskResponse};
use serde::Serialize;

/// Omni video generation module
///
/// This module provides comprehensive functionality for generating videos using the SkyReels Omni API,
/// supporting various reference types including keyframe images, subject references, scene references,
/// and video references for motion, editing, and composition tasks.
///
/// # Example
///
/// ```rust,no_run
/// use skyreels::{SkyReelsClient, Mode, ReferenceImageType, ReferenceVideoType};
///
/// # #[tokio::main]
/// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
/// let client = SkyReelsClient::new("your_api_key");
///
/// // Keyframe video generation
/// let request = client.omni()
///     .submit_request(
///         "The King @image1 anxiously searched for the destination on the map, \
///          suddenly made a surprise discovery, ran out of the door, and summoned a flying dragon. \
///          The dragon lowered its body @image2, the King mounted it, and the dragon took off. \
///          The King rode the dragon forward, flying over a village @image3."
///     )
///     .aspect_ratio("16:9")
///     .duration(5)
///     .add_reference_image(
///         "@image1",
///         ReferenceImageType::Keyframe,
///         vec!["https://example.com/frame1.jpg".to_string()],
///         Some(0)
///     )
///     .add_reference_image(
///         "@image2",
///         ReferenceImageType::Keyframe,
///         vec!["https://example.com/frame2.jpg".to_string()],
///         Some(2)
///     )
///     .add_reference_image(
///         "@image3",
///         ReferenceImageType::Keyframe,
///         vec!["https://example.com/frame3.jpg".to_string()],
///         Some(5)
///     )
///     .sound(false)
///     .prompt_optimizer(true)
///     .mode(Mode::Std)
///     .build();
///
/// let task_id = client.omni().submit(request).await?;
///
/// // Query task status
/// let status = client.omni()
///     .get_task(&task_id)
///     .await?;
///
/// # Ok(())
/// # }
/// ```

/// Type of image reference
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ReferenceImageType {
    /// Keyframe reference - displays a specific reference image at a designated time
    Keyframe,
    /// Subject reference - for character/subject consistency
    Subject,
    /// General image reference - for scene, style, or composition reference
    Image,
}

/// Type of video reference
#[derive(Debug, Clone, Serialize, PartialEq, Eq)]
#[serde(rename_all = "lowercase")]
pub enum ReferenceVideoType {
    /// Motion reference - imitate actions/movements from the reference video
    Reference,
    /// Base video - used for editing tasks (subject replacement, background replacement, etc.)
    Base,
}

/// Reference image configuration
#[derive(Debug, Clone, Serialize)]
pub struct ReferenceImage {
    /// Identifier for the reference group (must start with @ and appear in prompt)
    pub tag: String,
    /// Type of image reference
    #[serde(rename = "type")]
    pub ref_type: ReferenceImageType,
    /// List of image URLs
    pub image_urls: Vec<String>,
    /// Target timestamp for keyframe references (0 for first frame, duration for last)
    #[serde(skip_serializing_if = "Option::is_none")]
    pub time_stamp: Option<i32>,
}

impl ReferenceImage {
    /// Create a new reference image
    pub fn new(
        tag: impl Into<String>,
        ref_type: ReferenceImageType,
        image_urls: Vec<String>,
    ) -> Self {
        Self {
            tag: tag.into(),
            ref_type,
            image_urls,
            time_stamp: None,
        }
    }

    /// Set the timestamp for keyframe references
    pub fn time_stamp(mut self, time_stamp: i32) -> Self {
        self.time_stamp = Some(time_stamp);
        self
    }
}

/// Reference video configuration
#[derive(Debug, Clone, Serialize)]
pub struct ReferenceVideo {
    /// Identifier for the video reference (must start with @ and appear in prompt)
    pub tag: String,
    /// Type of video reference
    #[serde(rename = "type")]
    pub ref_type: ReferenceVideoType,
    /// The video URL (supports MP4, MOV, max 10s)
    pub video_url: String,
}

impl ReferenceVideo {
    /// Create a new reference video
    pub fn new(
        tag: impl Into<String>,
        ref_type: ReferenceVideoType,
        video_url: impl Into<String>,
    ) -> Self {
        Self {
            tag: tag.into(),
            ref_type,
            video_url: video_url.into(),
        }
    }
}

/// Builder for omni video task submission
#[derive(Debug, Clone, Serialize)]
pub struct OmniRequest {
    api_key: String,
    prompt: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    aspect_ratio: Option<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    duration: Option<i32>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    ref_images: Vec<ReferenceImage>,
    #[serde(skip_serializing_if = "Vec::is_empty")]
    ref_videos: Vec<ReferenceVideo>,
    #[serde(skip_serializing_if = "Option::is_none")]
    sound: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    prompt_optimizer: Option<bool>,
    #[serde(skip_serializing_if = "Option::is_none")]
    mode: Option<String>,
}

impl OmniRequest {
    /// Create a new omni request builder
    pub fn new(
        api_key: impl Into<String>,
        prompt: impl Into<String>,
    ) -> Self {
        Self {
            api_key: api_key.into(),
            prompt: prompt.into(),
            aspect_ratio: None,
            duration: None,
            ref_images: Vec::new(),
            ref_videos: Vec::new(),
            sound: None,
            prompt_optimizer: None,
            mode: None,
        }
    }

    /// Set the aspect ratio ('16:9', '4:3', '1:1', '9:16', '3:4')
    ///
    /// Note: This parameter is ignored if ref_videos is provided,
    /// as the output dimensions will automatically align with the reference video.
    pub fn aspect_ratio(mut self, aspect_ratio: impl Into<String>) -> Self {
        let ratio = aspect_ratio.into();
        let valid_ratios = ["16:9", "4:3", "1:1", "9:16", "3:4"];
        if !valid_ratios.contains(&ratio.as_str()) {
            tracing::warn!(
                "Aspect ratio '{}' is not valid. Valid values: {:?}. Using default '16:9'",
                ratio, valid_ratios
            );
            self.aspect_ratio = Some("16:9".to_string());
        } else {
            self.aspect_ratio = Some(ratio);
        }
        self
    }

    /// Set the video duration (3-15 seconds, default 5)
    ///
    /// Note: This parameter is overridden if ref_videos is provided;
    /// in such cases, the output duration will match the reference video (up to 10 seconds).
    pub fn duration(mut self, duration: i32) -> Self {
        if duration < 3 || duration > 15 {
            tracing::warn!("Duration {} is out of range [3, 15], clamping", duration);
            self.duration = Some(duration.clamp(3, 15));
        } else {
            self.duration = Some(duration);
        }
        self
    }

    /// Add a reference image
    ///
    /// # Parameters
    ///
    /// * `tag` - Identifier for the reference group (must start with @ and appear in prompt)
    /// * `ref_type` - Type of image reference
    /// * `image_urls` - List of image URLs
    /// * `time_stamp` - Target timestamp for keyframe references (required for Keyframe type)
    pub fn add_reference_image(
        mut self,
        tag: impl Into<String>,
        ref_type: ReferenceImageType,
        image_urls: Vec<String>,
        time_stamp: Option<i32>,
    ) -> Self {
        let tag_str = tag.into();
        if !tag_str.starts_with('@') {
            tracing::warn!("Reference tag '{}' should start with '@'", tag_str);
        }

        if ref_type == ReferenceImageType::Keyframe && time_stamp.is_none() {
            tracing::warn!(
                "Keyframe reference '{}' requires a time_stamp, but none was provided",
                tag_str
            );
        }

        let mut ref_image = ReferenceImage::new(tag_str, ref_type, image_urls);
        ref_image.time_stamp = time_stamp;
        self.ref_images.push(ref_image);
        self
    }

    /// Add a reference video
    ///
    /// # Parameters
    ///
    /// * `tag` - Identifier for the video reference (must start with @ and appear in prompt)
    /// * `ref_type` - Type of video reference
    /// * `video_url` - The video URL (supports MP4, MOV, max 10s)
    pub fn add_reference_video(
        mut self,
        tag: impl Into<String>,
        ref_type: ReferenceVideoType,
        video_url: impl Into<String>,
    ) -> Self {
        let tag_str = tag.into();
        if !tag_str.starts_with('@') {
            tracing::warn!("Reference tag '{}' should start with '@'", tag_str);
        }

        self.ref_videos.push(ReferenceVideo::new(tag_str, ref_type, video_url));
        self
    }

    /// Set whether to include sound effects (default false)
    ///
    /// Note: When using video reference, sound is not effective; no audio by default.
    pub fn sound(mut self, sound: bool) -> Self {
        self.sound = Some(sound);
        self
    }

    /// Set whether to enable prompt optimizer (default true)
    ///
    /// Enables automatic prompt expansion and refinement to achieve higher visual fidelity and alignment.
    pub fn prompt_optimizer(mut self, prompt_optimizer: bool) -> Self {
        self.prompt_optimizer = Some(prompt_optimizer);
        self
    }

    /// Set the quality/performance mode
    ///
    /// * `fast` - Faster generation (currently not supported, will be added later)
    /// * `std` - Balanced speed and quality (default)
    /// * `pro` - Higher quality (currently not supported, will be added later)
    ///
    /// All modes output at 1080p.
    pub fn mode(mut self, mode: Mode) -> Self {
        self.mode = Some(mode.to_string());
        self
    }

    /// Build the request
    pub fn build(self) -> Self {
        // Apply defaults
        let mut request = self;

        if request.aspect_ratio.is_none() {
            request.aspect_ratio = Some("16:9".to_string());
        }

        if request.duration.is_none() {
            request.duration = Some(5);
        }

        if request.sound.is_none() {
            request.sound = Some(false);
        }

        if request.prompt_optimizer.is_none() {
            request.prompt_optimizer = Some(true);
        }

        if request.mode.is_none() {
            request.mode = Some("std".to_string());
        }

        request
    }
}

/// Omni video API handler
pub struct Omni<'a> {
    pub(crate) client: &'a SkyReelsClient,
}

impl<'a> Omni<'a> {
    /// Create a new omni handler
    pub(crate) fn new(client: &'a SkyReelsClient) -> Self {
        Self { client }
    }

    /// Create a new omni request builder
    pub fn submit_request(&self, prompt: impl Into<String>) -> OmniRequest {
        OmniRequest::new(self.client.api_key_sync(), prompt)
    }

    /// Submit an omni video task
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use skyreels::{SkyReelsClient, ReferenceImageType, ReferenceVideoType, Mode};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = SkyReelsClient::new("your_api_key");
    ///
    /// // Video motion reference
    /// let request = client.omni()
    ///     .submit_request(
    ///         "The man in @image_1 imitates the actions of the red male in @video_1, \
    ///          and the monkey in @image_2 imitates the actions of the dog in @video_1, \
    ///          while keeping the background of @video_1."
    ///     )
    ///     .add_reference_image(
    ///         "@image_1",
    ///         ReferenceImageType::Image,
    ///         vec!["https://example.com/man.jpg".to_string()],
    ///         None
    ///     )
    ///     .add_reference_image(
    ///         "@image_2",
    ///         ReferenceImageType::Image,
    ///         vec!["https://example.com/monkey.jpg".to_string()],
    ///         None
    ///     )
    ///     .add_reference_video(
    ///         "@video_1",
    ///         ReferenceVideoType::Reference,
    ///         "https://example.com/reference.mp4"
    ///     )
    ///     .sound(false)
    ///     .mode(Mode::Std)
    ///     .build();
    ///
    /// let task_id = client.omni().submit(request).await?;
    /// # Ok(())
    /// # }
    /// ```
    pub async fn submit(&self, request: OmniRequest) -> Result<String> {
        let base_url = self.client.base_url_sync();
        let url = format!("{}/video/omni-video/submit", base_url);

        tracing::debug!("Submitting omni video task to: {}", url);

        let response = self.client
            .http_client()
            .post(&url)
            .json(&request)
            .send()
            .await?;

        let status = response.status();
        let response_text = response.text().await?;

        // Output raw JSON response for debugging
        tracing::debug!("Raw API response (submit): {}", response_text);

        if !status.is_success() {
            return Err(Error::api_error(
                status.as_u16(),
                format!("{}: {}", status.as_u16(), response_text),
            ));
        }

        let submit_response: SubmitResponse = serde_json::from_str(&response_text)
            .map_err(|e| Error::ResponseParseError(format!("Failed to parse response: {}", e)))?;

        // Log parsed response for debugging
        tracing::debug!("Parsed submit response: task_id={}, code={}, status={}",
            submit_response.task_id, submit_response.code, submit_response.status);

        if submit_response.code != 200 {
            return Err(Error::SubmissionFailed(submit_response.msg));
        }

        tracing::info!("Omni task submitted successfully: {}", submit_response.task_id);

        Ok(submit_response.task_id)
    }

    /// Query task status by task ID
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use skyreels::SkyReelsClient;
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = SkyReelsClient::new("your_api_key");
    ///
    /// let status = client.omni()
    ///     .get_task("task_id_here")
    ///     .await?;
    ///
    /// println!("Status: {}", status.get_status());
    /// if let Some(video) = status.get_video_data() {
    ///     println!("Video URL: {}", video.video_url);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn get_task(&self, task_id: &str) -> Result<TaskResponse> {
        let base_url = self.client.base_url_sync();
        let url = format!("{}/video/omni-video/task/{}", base_url, task_id);

        tracing::debug!("Querying omni task status: {}", url);

        let response = self.client
            .http_client()
            .get(&url)
            .send()
            .await?;

        let status = response.status();
        let response_text = response.text().await?;

        // Output raw JSON response for debugging
        tracing::debug!("Raw API response (query): {}", response_text);

        if !status.is_success() {
            return Err(Error::api_error(
                status.as_u16(),
                format!("{}: {}", status.as_u16(), response_text),
            ));
        }

        let task_response: TaskResponse = serde_json::from_str(&response_text)
            .map_err(|e| Error::ResponseParseError(format!("Failed to parse response: {}", e)))?;

        // Log parsed response for debugging
        tracing::debug!("Parsed task response: task_id={}, status={}, code={}",
            task_response.task_id, task_response.status, task_response.code);

        if let Some(ref video_data) = task_response.data {
            tracing::debug!("Video data: url={}, duration={}, resolution={}",
                video_data.video_url, video_data.duration, video_data.resolution);
        }

        tracing::info!("Task {} status: {}", task_id, task_response.get_status());

        Ok(task_response)
    }

    /// Wait for task completion (polling)
    ///
    /// This method polls the task status until it completes or fails.
    ///
    /// # Parameters
    ///
    /// * `task_id` - The task ID to wait for
    /// * `poll_interval_secs` - Polling interval in seconds (default: 5)
    /// * `max_attempts` - Maximum number of polling attempts (default: 120)
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// # use skyreels::{SkyReelsClient, ReferenceImageType, Mode};
    /// # #[tokio::main]
    /// # async fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let client = SkyReelsClient::new("your_api_key");
    ///
    /// let request = client.omni()
    ///     .submit_request("A beautiful video with @image1")
    ///     .add_reference_image(
    ///         "@image1",
    ///         ReferenceImageType::Image,
    ///         vec!["https://example.com/image.jpg".to_string()],
    ///         None
    ///     )
    ///     .build();
    ///
    /// let task_id = client.omni().submit(request).await?;
    ///
    /// let result = client.omni()
    ///     .wait_for_completion(&task_id, 5, 120)
    ///     .await?;
    ///
    /// if let Some(video) = result.get_video_data() {
    ///     println!("Video ready: {}", video.video_url);
    /// }
    /// # Ok(())
    /// # }
    /// ```
    pub async fn wait_for_completion(
        &self,
        task_id: &str,
        poll_interval_secs: u64,
        max_attempts: usize,
    ) -> Result<TaskResponse> {
        let mut attempts = 0;

        loop {
            attempts += 1;

            let response = self.get_task(task_id).await?;
            let status = response.get_status();

            if status.is_completed() {
                if status == TaskStatus::Success {
                    tracing::info!("Task {} completed successfully", task_id);
                    return Ok(response);
                } else {
                    return Err(Error::QueryFailed(format!(
                        "Task {} failed: {}",
                        task_id, response.msg
                    )));
                }
            }

            if attempts >= max_attempts {
                return Err(Error::QueryFailed(format!(
                    "Task {} did not complete after {} attempts",
                    task_id, max_attempts
                )));
            }

            tracing::debug!(
                "Task {} is {}, waiting {} seconds... (attempt {}/{})",
                task_id,
                status,
                poll_interval_secs,
                attempts,
                max_attempts
            );

            tokio::time::sleep(tokio::time::Duration::from_secs(poll_interval_secs)).await;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_request_creation() {
        let request = OmniRequest::new("test_key", "A beautiful video");
        assert_eq!(request.prompt, "A beautiful video");
        assert!(request.ref_images.is_empty());
        assert!(request.ref_videos.is_empty());
    }

    #[test]
    fn test_aspect_ratio_validation() {
        let request = OmniRequest::new("test_key", "test")
            .aspect_ratio("9:16")
            .build();
        assert_eq!(request.aspect_ratio, Some("9:16".to_string()));

        let request2 = OmniRequest::new("test_key", "test")
            .aspect_ratio("invalid")
            .build();
        assert_eq!(request2.aspect_ratio, Some("16:9".to_string()));
    }

    #[test]
    fn test_duration_clamping() {
        let request = OmniRequest::new("test_key", "test")
            .duration(20)
            .build();
        assert_eq!(request.duration, Some(15));

        let request2 = OmniRequest::new("test_key", "test")
            .duration(1)
            .build();
        assert_eq!(request2.duration, Some(3));
    }

    #[test]
    fn test_reference_image() {
        let ref_image = ReferenceImage::new(
            "@image1",
            ReferenceImageType::Keyframe,
            vec!["https://example.com/image.jpg".to_string()],
        ).time_stamp(0);

        assert_eq!(ref_image.tag, "@image1");
        assert_eq!(ref_image.ref_type, ReferenceImageType::Keyframe);
        assert_eq!(ref_image.image_urls.len(), 1);
        assert_eq!(ref_image.time_stamp, Some(0));
    }

    #[test]
    fn test_reference_video() {
        let ref_video = ReferenceVideo::new(
            "@video1",
            ReferenceVideoType::Reference,
            "https://example.com/video.mp4"
        );

        assert_eq!(ref_video.tag, "@video1");
        assert_eq!(ref_video.ref_type, ReferenceVideoType::Reference);
        assert_eq!(ref_video.video_url, "https://example.com/video.mp4");
    }

    #[test]
    fn test_request_with_references() {
        let request = OmniRequest::new("test_key", "Test @image1 and @video1")
            .add_reference_image(
                "@image1",
                ReferenceImageType::Image,
                vec!["https://example.com/image.jpg".to_string()],
                None
            )
            .add_reference_video(
                "@video1",
                ReferenceVideoType::Reference,
                "https://example.com/video.mp4"
            )
            .build();

        assert_eq!(request.ref_images.len(), 1);
        assert_eq!(request.ref_videos.len(), 1);
        assert_eq!(request.ref_images[0].tag, "@image1");
        assert_eq!(request.ref_videos[0].tag, "@video1");
    }

    #[test]
    fn test_defaults() {
        let request = OmniRequest::new("test_key", "test").build();

        assert_eq!(request.aspect_ratio, Some("16:9".to_string()));
        assert_eq!(request.duration, Some(5));
        assert_eq!(request.sound, Some(false));
        assert_eq!(request.prompt_optimizer, Some(true));
        assert_eq!(request.mode, Some("std".to_string()));
    }
}
