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, pub context: Option, pub repo_path: Option, pub fallback_model: Option, } #[derive(Deserialize)] pub struct FixCodeRequest { pub code: String, pub error_message: Option, pub language: Option, } #[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, } #[derive(Serialize)] pub struct CodeResponse { pub code: String, pub explanation: String, pub language: String, pub detected_stack: String, pub files: Vec, pub workspace_artifact_dir: Option, pub bundle_path: Option, pub bundle_download_url: Option, pub entrypoint: Option, pub repo_path: Option, } #[derive(Serialize)] pub struct CodeDiffFileResponse { pub path: String, pub status: String, pub patch: Option, } #[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, } #[derive(Serialize)] pub struct CodeWorkspaceMutationResponse { pub artifact_dir: String, pub repo_path: String, pub applied_files: Vec, pub imported_from_bundle: bool, } fn bundle_download_url(workspace_artifact_dir: Option<&str>) -> Option { 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, ) -> (StatusCode, Json) { (status, Json(serde_json::json!({ "error": message.into() }))) } fn validate_artifact_dir( artifact_dir: &str, ) -> Result)> { 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 { 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)> { 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>>, 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> { 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) -> 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 { 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)> { 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, ¤t, &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, (StatusCode, Json)> { 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, (StatusCode, Json)> { 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, ) -> Result, (StatusCode, Json)> { 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::("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, ) -> Result, (StatusCode, Json)> { 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::("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, ) -> Result)> { 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, ) -> Result, (StatusCode, Json)> { 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, ) -> Result, (StatusCode, Json)> { 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, ) -> Result, (StatusCode, Json)> { 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, })) }