use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use std::time::Duration;
use tokio::process::Command;
use tokio::time::timeout;

use super::executor::PythonExecutor;

const ASRX_TIMEOUT: Duration = Duration::from_secs(7200);

#[derive(Debug, Serialize, Deserialize)]
pub struct AsrxResult {
    pub language: Option<String>,
    pub segments: Vec<AsrxSegment>,
}

#[derive(Debug, Serialize, Deserialize)]
pub struct AsrxSegment {
    pub start: f64,
    pub end: f64,
    pub text: String,
    pub speaker_id: Option<String>,
}

pub async fn process_asrx(
    video_path: &str,
    output_path: &str,
    uuid: Option<&str>,
) -> Result<AsrxResult> {
    let executor = PythonExecutor::new()?;
    let script_path = executor.script_path("asrx_processor_custom.py");

    tracing::info!(
        "[ASRX] Starting speaker diarization (custom): {}",
        video_path
    );

    if !script_path.exists() {
        tracing::warn!("[ASRX] Custom script not found, falling back to original");
        let fallback_path = executor.script_path("asrx_processor.py");
        if !fallback_path.exists() {
            tracing::warn!("[ASRX] No script found, returning empty result");
            return Ok(AsrxResult {
                language: None,
                segments: vec![],
            });
        }
    }

    let mut cmd = Command::new(executor.python_path());
    cmd.arg(&script_path).arg(video_path).arg(output_path);

    if let Some(u) = uuid {
        cmd.arg("--uuid").arg(u);
    }

    cmd.stdout(std::process::Stdio::piped())
        .stderr(std::process::Stdio::piped());

    let child = cmd.spawn().context("Failed to run ASRX processor")?;

    let output = match timeout(ASRX_TIMEOUT, child.wait_with_output()).await {
        Ok(Ok(output)) => output,
        Ok(Err(e)) => return Err(e).context("Failed to run ASRX processor"),
        Err(_) => anyhow::bail!("ASRX processing timed out after {:?}", ASRX_TIMEOUT),
    };

    let stderr = String::from_utf8_lossy(&output.stderr);

    for line in stderr.lines() {
        if line.starts_with("ASRX_START") {
            tracing::info!("[ASRX] Loading model...");
        } else if line.starts_with("ASRX_PROGRESS:") {
            let count = line.trim_start_matches("ASRX_PROGRESS:");
            tracing::info!("[ASRX] Processed {} segments...", count);
        } else if line.starts_with("ASRX_COMPLETE:") {
            let count = line.trim_start_matches("ASRX_COMPLETE:");
            tracing::info!("[ASRX] Completed! Total: {} segments", count);
        }
    }

    if !output.status.success() {
        anyhow::bail!("ASRX failed: {}", stderr);
    }

    let json_str = std::fs::read_to_string(output_path).context("Failed to read ASRX output")?;

    let result: AsrxResult =
        serde_json::from_str(&json_str).context("Failed to parse ASRX output")?;

    tracing::info!("[ASRX] Result: {} segments", result.segments.len());

    Ok(result)
}

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

    #[test]
    fn test_asrx_result_serialization() {
        let result = AsrxResult {
            language: Some("en".to_string()),
            segments: vec![AsrxSegment {
                start: 0.0,
                end: 2.5,
                text: "Hello".to_string(),
                speaker_id: Some("SPEAKER_00".to_string()),
            }],
        };

        let json = serde_json::to_string(&result).unwrap();
        assert!(json.contains("Hello"));
        assert!(json.contains("SPEAKER_00"));
    }

    #[test]
    fn test_asrx_result_deserialization() {
        let json = r#"{
            "language": "zh",
            "segments": [
                {"start": 0.0, "end": 1.5, "text": "測試", "speaker_id": "SPEAKER_01"}
            ]
        }"#;

        let result: AsrxResult = serde_json::from_str(json).unwrap();
        assert_eq!(result.language, Some("zh".to_string()));
        assert_eq!(result.segments.len(), 1);
        assert_eq!(
            result.segments[0].speaker_id,
            Some("SPEAKER_01".to_string())
        );
    }

    #[test]
    fn test_asrx_result_empty_segments() {
        let result = AsrxResult {
            language: None,
            segments: vec![],
        };
        assert!(result.segments.is_empty());
        assert!(result.language.is_none());
    }

    #[test]
    fn test_asrx_segment_times() {
        let segment = AsrxSegment {
            start: 0.0,
            end: 5.0,
            text: "Test".to_string(),
            speaker_id: None,
        };
        assert!(segment.end > segment.start);
    }
}
