backend-rust/src/api/code.rs

use std::{
    fs::{self, File},
    io::Cursor,
    path::{Path, PathBuf},
};

use axum::{
    extract::Query,
    http::{header, HeaderMap, HeaderValue, StatusCode},
    response::IntoResponse,
    Json,
};
use serde::{Deserialize, Serialize};
use zip::{write::SimpleFileOptions, CompressionMethod, ZipArchive, ZipWriter};

use crate::inference::python_bridge::PythonBridge;
use crate::inference::response_validator::CodeGenerationPayload;

const GENERATED_PROJECTS_ROOT: &str = "/tmp/maris-ai/generated-projects";

#[derive(Deserialize)]
pub struct CodeRequest {
    pub prompt: String,
    pub language: Option<String>,
    pub context: Option<String>,
    pub repo_path: Option<String>,
    pub fallback_model: Option<String>,
}

#[derive(Deserialize)]
pub struct FixCodeRequest {
    pub code: String,
    pub error_message: Option<String>,
    pub language: Option<String>,
}

#[derive(Deserialize)]
pub struct DownloadProjectZipRequest {
    pub artifact_dir: String,
}

#[derive(Deserialize)]
pub struct CodeWorkspaceQuery {
    pub artifact_dir: String,
    pub repo_path: String,
}

#[derive(Deserialize)]
pub struct CodeWorkspaceMutationRequest {
    pub artifact_dir: String,
    pub repo_path: String,
}

#[derive(Serialize)]
pub struct CodeProjectFileResponse {
    pub path: String,
    pub absolute_path: Option<String>,
}

#[derive(Serialize)]
pub struct CodeResponse {
    pub code: String,
    pub explanation: String,
    pub language: String,
    pub detected_stack: String,
    pub files: Vec<CodeProjectFileResponse>,
    pub workspace_artifact_dir: Option<String>,
    pub bundle_path: Option<String>,
    pub bundle_download_url: Option<String>,
    pub entrypoint: Option<String>,
    pub repo_path: Option<String>,
}

#[derive(Serialize)]
pub struct CodeDiffFileResponse {
    pub path: String,
    pub status: String,
    pub patch: Option<String>,
}

#[derive(Serialize)]
pub struct CodePreviewResponse {
    pub artifact_dir: String,
    pub repo_path: String,
    pub has_changes: bool,
    pub added_files: usize,
    pub modified_files: usize,
    pub unchanged_files: usize,
    pub patch: String,
    pub files: Vec<CodeDiffFileResponse>,
}

#[derive(Serialize)]
pub struct CodeWorkspaceMutationResponse {
    pub artifact_dir: String,
    pub repo_path: String,
    pub applied_files: Vec<String>,
    pub imported_from_bundle: bool,
}

fn bundle_download_url(workspace_artifact_dir: Option<&str>) -> Option<String> {
    workspace_artifact_dir.map(|artifact_dir| {
        format!(
            "/api/code/download?artifact_dir={}",
            urlencoding::encode(artifact_dir)
        )
    })
}

fn generated_projects_root() -> PathBuf {
    PathBuf::from(GENERATED_PROJECTS_ROOT)
}

fn json_error(
    status: StatusCode,
    message: impl Into<String>,
) -> (StatusCode, Json<serde_json::Value>) {
    (status, Json(serde_json::json!({ "error": message.into() })))
}

fn validate_artifact_dir(
    artifact_dir: &str,
) -> Result<PathBuf, (StatusCode, Json<serde_json::Value>)> {
    let trimmed = artifact_dir.trim();
    if trimmed.is_empty() {
        return Err(json_error(
            StatusCode::BAD_REQUEST,
            "Trūkst artifact_dir parametra.",
        ));
    }

    let requested = PathBuf::from(trimmed);
    let canonical_requested = requested.canonicalize().map_err(|_| {
        json_error(
            StatusCode::NOT_FOUND,
            "Norādītā artefakta direktorija nav atrasta.",
        )
    })?;
    let root = generated_projects_root();
    let canonical_root = root.canonicalize().unwrap_or(root);
    if !canonical_requested.starts_with(&canonical_root) || !canonical_requested.is_dir() {
        return Err(json_error(
            StatusCode::BAD_REQUEST,
            "Nederīga artefakta direktorija.",
        ));
    }
    Ok(canonical_requested)
}

fn find_repo_root(path: &Path) -> Option<PathBuf> {
    let mut current = Some(path);
    while let Some(candidate) = current {
        if candidate.join(".git").exists() {
            return Some(candidate.to_path_buf());
        }
        current = candidate.parent();
    }
    None
}

fn validate_repo_path(repo_path: &str) -> Result<PathBuf, (StatusCode, Json<serde_json::Value>)> {
    let trimmed = repo_path.trim();
    if trimmed.is_empty() {
        return Err(json_error(
            StatusCode::BAD_REQUEST,
            "Trūkst repo_path parametra.",
        ));
    }

    let requested = PathBuf::from(trimmed);
    let canonical_requested = requested.canonicalize().map_err(|_| {
        json_error(
            StatusCode::NOT_FOUND,
            "Norādītā repozitorija direktorija nav atrasta.",
        )
    })?;
    let candidate_dir = if canonical_requested.is_file() {
        canonical_requested
            .parent()
            .map(Path::to_path_buf)
            .ok_or_else(|| {
                json_error(
                    StatusCode::BAD_REQUEST,
                    "Norādītais repo_path nav derīga direktorija.",
                )
            })?
    } else {
        canonical_requested
    };

    let repo_root = find_repo_root(&candidate_dir).ok_or_else(|| {
        json_error(
            StatusCode::BAD_REQUEST,
            "repo_path nav Git repozitorija darba telpā.",
        )
    })?;
    if !candidate_dir.starts_with(&repo_root) {
        return Err(json_error(
            StatusCode::BAD_REQUEST,
            "repo_path atrodas ārpus Git repozitorija saknes.",
        ));
    }

    Ok(candidate_dir)
}

fn append_directory_to_zip(
    writer: &mut ZipWriter<Cursor<Vec<u8>>>,
    root: &Path,
    current: &Path,
) -> anyhow::Result<()> {
    let options = SimpleFileOptions::default().compression_method(CompressionMethod::Deflated);
    for entry in fs::read_dir(current)? {
        let entry = entry?;
        let path = entry.path();
        let relative = path
            .strip_prefix(root)?
            .to_string_lossy()
            .replace('\\', "/");
        if path.is_dir() {
            if !relative.is_empty() {
                writer.add_directory(format!("{relative}/"), options)?;
            }
            append_directory_to_zip(writer, root, &path)?;
        } else if path.is_file() {
            writer.start_file(relative, options)?;
            let mut file = File::open(&path)?;
            std::io::copy(&mut file, writer)?;
        }
    }
    Ok(())
}

fn zip_directory_to_bytes(artifact_dir: &Path) -> anyhow::Result<Vec<u8>> {
    let cursor = Cursor::new(Vec::new());
    let mut writer = ZipWriter::new(cursor);
    append_directory_to_zip(&mut writer, artifact_dir, artifact_dir)?;
    let cursor = writer.finish()?;
    Ok(cursor.into_inner())
}

fn collect_files(root: &Path, current: &Path, files: &mut Vec<PathBuf>) -> anyhow::Result<()> {
    for entry in fs::read_dir(current)? {
        let entry = entry?;
        let path = entry.path();
        if path.is_dir() {
            collect_files(root, &path, files)?;
        } else if path.is_file() {
            files.push(path.strip_prefix(root)?.to_path_buf());
        }
    }
    Ok(())
}

fn read_text_lossy(path: &Path) -> anyhow::Result<String> {
    Ok(String::from_utf8_lossy(&fs::read(path)?).into_owned())
}

fn build_patch(relative_path: &str, current_content: &str, generated_content: &str) -> String {
    let current_lines: Vec<&str> = current_content.lines().collect();
    let generated_lines: Vec<&str> = generated_content.lines().collect();
    let current_start = if current_lines.is_empty() { 0 } else { 1 };
    let generated_start = if generated_lines.is_empty() { 0 } else { 1 };
    let mut patch = format!(
        "--- a/{relative_path}\n+++ b/{relative_path}\n@@ -{current_start},{} +{generated_start},{} @@\n",
        current_lines.len(),
        generated_lines.len()
    );
    for line in current_lines {
        patch.push('-');
        patch.push_str(line);
        patch.push('\n');
    }
    for line in generated_lines {
        patch.push('+');
        patch.push_str(line);
        patch.push('\n');
    }
    patch
}

fn build_preview_response(
    artifact_dir: &Path,
    repo_path: &Path,
) -> Result<CodePreviewResponse, (StatusCode, Json<serde_json::Value>)> {
    let mut relative_files = Vec::new();
    collect_files(artifact_dir, artifact_dir, &mut relative_files).map_err(|err| {
        json_error(
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Neizdevās nolasīt workspace failus: {err}"),
        )
    })?;
    relative_files.sort();

    let mut files = Vec::new();
    let mut combined_patch = String::new();
    let mut added_files = 0usize;
    let mut modified_files = 0usize;
    let mut unchanged_files = 0usize;

    for relative in relative_files {
        let relative_display = relative.to_string_lossy().replace('\\', "/");
        let generated_path = artifact_dir.join(&relative);
        let generated_content = read_text_lossy(&generated_path).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās nolasīt ģenerēto failu `{relative_display}`: {err}"),
            )
        })?;
        let repo_file = repo_path.join(&relative);
        let current_content = if repo_file.exists() {
            Some(read_text_lossy(&repo_file).map_err(|err| {
                json_error(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Neizdevās nolasīt repo failu `{relative_display}`: {err}"),
                )
            })?)
        } else {
            None
        };

        let (status, patch) = match current_content {
            None => {
                added_files += 1;
                (
                    "added".to_string(),
                    Some(build_patch(&relative_display, "", &generated_content)),
                )
            }
            Some(current) if current == generated_content => {
                unchanged_files += 1;
                ("unchanged".to_string(), None)
            }
            Some(current) => {
                modified_files += 1;
                (
                    "modified".to_string(),
                    Some(build_patch(&relative_display, &current, &generated_content)),
                )
            }
        };

        if let Some(ref patch_body) = patch {
            if !combined_patch.is_empty() {
                combined_patch.push_str("\n\n");
            }
            combined_patch.push_str(patch_body);
        }

        files.push(CodeDiffFileResponse {
            path: relative_display,
            status,
            patch,
        });
    }

    Ok(CodePreviewResponse {
        artifact_dir: artifact_dir.to_string_lossy().into_owned(),
        repo_path: repo_path.to_string_lossy().into_owned(),
        has_changes: added_files > 0 || modified_files > 0,
        added_files,
        modified_files,
        unchanged_files,
        patch: combined_patch,
        files,
    })
}

fn copy_artifact_files_to_repo(
    artifact_dir: &Path,
    repo_path: &Path,
) -> Result<Vec<String>, (StatusCode, Json<serde_json::Value>)> {
    let mut relative_files = Vec::new();
    collect_files(artifact_dir, artifact_dir, &mut relative_files).map_err(|err| {
        json_error(
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Neizdevās nolasīt workspace failus: {err}"),
        )
    })?;
    relative_files.sort();

    let mut applied_files = Vec::new();
    for relative in relative_files {
        let relative_display = relative.to_string_lossy().replace('\\', "/");
        let source = artifact_dir.join(&relative);
        let target = repo_path.join(&relative);
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent).map_err(|err| {
                json_error(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Neizdevās sagatavot direktoriju `{relative_display}`: {err}"),
                )
            })?;
        }
        fs::copy(&source, &target).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās ierakstīt failu `{relative_display}` repozitorijā: {err}"),
            )
        })?;
        applied_files.push(relative_display);
    }
    Ok(applied_files)
}

fn import_bundle_into_repo(
    artifact_dir: &Path,
    repo_path: &Path,
) -> Result<Vec<String>, (StatusCode, Json<serde_json::Value>)> {
    let bundle_path = artifact_dir.with_extension("zip");
    let bundle_bytes = if bundle_path.exists() {
        fs::read(&bundle_path).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās nolasīt zip bundle: {err}"),
            )
        })?
    } else {
        zip_directory_to_bytes(artifact_dir).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās izveidot zip bundle: {err}"),
            )
        })?
    };

    let mut archive = ZipArchive::new(Cursor::new(bundle_bytes)).map_err(|err| {
        json_error(
            StatusCode::INTERNAL_SERVER_ERROR,
            format!("Neizdevās atvērt zip bundle: {err}"),
        )
    })?;
    let mut applied_files = Vec::new();
    for index in 0..archive.len() {
        let mut entry = archive.by_index(index).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās nolasīt zip ierakstu: {err}"),
            )
        })?;
        let Some(relative) = entry.enclosed_name() else {
            return Err(json_error(
                StatusCode::BAD_REQUEST,
                "Zip bundle satur nederīgu ceļu.",
            ));
        };
        let relative_display = relative.to_string_lossy().replace('\\', "/");
        let target = repo_path.join(relative);
        if entry.is_dir() {
            fs::create_dir_all(&target).map_err(|err| {
                json_error(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Neizdevās izveidot direktoriju `{relative_display}`: {err}"),
                )
            })?;
            continue;
        }
        if let Some(parent) = target.parent() {
            fs::create_dir_all(parent).map_err(|err| {
                json_error(
                    StatusCode::INTERNAL_SERVER_ERROR,
                    format!("Neizdevās sagatavot direktoriju `{relative_display}`: {err}"),
                )
            })?;
        }
        let mut file = File::create(&target).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās sagatavot failu `{relative_display}`: {err}"),
            )
        })?;
        std::io::copy(&mut entry, &mut file).map_err(|err| {
            json_error(
                StatusCode::INTERNAL_SERVER_ERROR,
                format!("Neizdevās importēt failu `{relative_display}`: {err}"),
            )
        })?;
        applied_files.push(relative_display);
    }
    applied_files.sort();
    Ok(applied_files)
}

pub async fn generate_code(
    Json(req): Json<CodeRequest>,
) -> Result<Json<CodeResponse>, (StatusCode, Json<serde_json::Value>)> {
    let bridge = PythonBridge::new();
    let payload = serde_json::json!({
        "prompt": req.prompt,
        "language": req.language.clone().unwrap_or_else(|| "Python".to_string()),
        "context": req.context.unwrap_or_default(),
        "repo_path": req.repo_path,
        "fallback_model": req.fallback_model,
    });

    match bridge
        .call::<CodeGenerationPayload>("code/generate", &payload)
        .await
    {
        Ok(result) => Ok(Json(CodeResponse {
            code: result.code,
            explanation: result.explanation,
            language: result.language,
            detected_stack: result.detected_stack,
            files: result
                .files
                .into_iter()
                .map(|file| CodeProjectFileResponse {
                    path: file.path,
                    absolute_path: file.absolute_path,
                })
                .collect(),
            workspace_artifact_dir: result.workspace_artifact_dir.clone(),
            bundle_path: result.bundle_path,
            bundle_download_url: bundle_download_url(result.workspace_artifact_dir.as_deref()),
            entrypoint: result.entrypoint,
            repo_path: result.repo_path,
        })),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": e.to_string() })),
        )),
    }
}

pub async fn fix_code(
    Json(req): Json<FixCodeRequest>,
) -> Result<Json<CodeResponse>, (StatusCode, Json<serde_json::Value>)> {
    let bridge = PythonBridge::new();
    let payload = serde_json::json!({
        "code": req.code,
        "error_message": req.error_message.unwrap_or_default(),
        "language": req.language.clone().unwrap_or_else(|| "Python".to_string()),
    });

    match bridge
        .call::<CodeGenerationPayload>("code/fix", &payload)
        .await
    {
        Ok(result) => Ok(Json(CodeResponse {
            code: result.code,
            explanation: result.explanation,
            language: result.language,
            detected_stack: result.detected_stack,
            files: result
                .files
                .into_iter()
                .map(|file| CodeProjectFileResponse {
                    path: file.path,
                    absolute_path: file.absolute_path,
                })
                .collect(),
            workspace_artifact_dir: result.workspace_artifact_dir.clone(),
            bundle_path: result.bundle_path,
            bundle_download_url: bundle_download_url(result.workspace_artifact_dir.as_deref()),
            entrypoint: result.entrypoint,
            repo_path: result.repo_path,
        })),
        Err(e) => Err((
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({ "error": e.to_string() })),
        )),
    }
}

pub async fn download_project_zip(
    Query(req): Query<DownloadProjectZipRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<serde_json::Value>)> {
    let artifact_dir = validate_artifact_dir(&req.artifact_dir)?;
    let bundle_path = artifact_dir.with_extension("zip");
    let bytes = if bundle_path.exists() {
        fs::read(&bundle_path).map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": format!("Neizdevās nolasīt bundle: {e}") })),
            )
        })?
    } else {
        zip_directory_to_bytes(&artifact_dir).map_err(|e| {
            (
                StatusCode::INTERNAL_SERVER_ERROR,
                Json(serde_json::json!({ "error": format!("Neizdevās izveidot zip bundle: {e}") })),
            )
        })?
    };

    let file_stem = artifact_dir
        .file_name()
        .and_then(|name| name.to_str())
        .unwrap_or("project-bundle");
    let mut headers = HeaderMap::new();
    headers.insert(
        header::CONTENT_TYPE,
        HeaderValue::from_static("application/zip"),
    );
    headers.insert(
        header::CONTENT_DISPOSITION,
        HeaderValue::from_str(&format!("attachment; filename=\"{file_stem}.zip\""))
            .unwrap_or_else(|_| HeaderValue::from_static("attachment")),
    );
    headers.insert(
        header::CACHE_CONTROL,
        HeaderValue::from_static("no-store, max-age=0"),
    );

    Ok((headers, bytes))
}

pub async fn preview_code_changes(
    Query(req): Query<CodeWorkspaceQuery>,
) -> Result<Json<CodePreviewResponse>, (StatusCode, Json<serde_json::Value>)> {
    let artifact_dir = validate_artifact_dir(&req.artifact_dir)?;
    let repo_path = validate_repo_path(&req.repo_path)?;
    Ok(Json(build_preview_response(&artifact_dir, &repo_path)?))
}

pub async fn apply_workspace_to_repo(
    Json(req): Json<CodeWorkspaceMutationRequest>,
) -> Result<Json<CodeWorkspaceMutationResponse>, (StatusCode, Json<serde_json::Value>)> {
    let artifact_dir = validate_artifact_dir(&req.artifact_dir)?;
    let repo_path = validate_repo_path(&req.repo_path)?;
    let applied_files = copy_artifact_files_to_repo(&artifact_dir, &repo_path)?;
    Ok(Json(CodeWorkspaceMutationResponse {
        artifact_dir: artifact_dir.to_string_lossy().into_owned(),
        repo_path: repo_path.to_string_lossy().into_owned(),
        applied_files,
        imported_from_bundle: false,
    }))
}

pub async fn import_bundle_to_repo(
    Json(req): Json<CodeWorkspaceMutationRequest>,
) -> Result<Json<CodeWorkspaceMutationResponse>, (StatusCode, Json<serde_json::Value>)> {
    let artifact_dir = validate_artifact_dir(&req.artifact_dir)?;
    let repo_path = validate_repo_path(&req.repo_path)?;
    let applied_files = import_bundle_into_repo(&artifact_dir, &repo_path)?;
    Ok(Json(CodeWorkspaceMutationResponse {
        artifact_dir: artifact_dir.to_string_lossy().into_owned(),
        repo_path: repo_path.to_string_lossy().into_owned(),
        applied_files,
        imported_from_bundle: true,
    }))
}