use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{Read, Write};
use std::path::{Path, PathBuf};

use clap::ValueEnum;
use flate2::write::GzEncoder;
use flate2::Compression;
use serde::Serialize;
use walkdir::WalkDir;
use zip::write::SimpleFileOptions;

use crate::archive::ArchiveLayout;
use crate::config::Config;
use crate::errors::{Result, SiteforgeError};
use crate::parse::{estimate_tokens, JsonlChunk};
use crate::storage::Storage;

#[derive(Debug, Clone, Copy, PartialEq, Eq, ValueEnum)]
pub enum ExportFormat {
    Markdown,
    BasicMarkdown,
    Jsonl,
    Json,
    Html,
    Warc,
    Zip,
    Tar,
    AgentTar,
    All,
}

impl ExportFormat {
    pub fn as_str(self) -> &'static str {
        match self {
            ExportFormat::Markdown => "markdown",
            ExportFormat::BasicMarkdown => "basic-markdown",
            ExportFormat::Jsonl => "jsonl",
            ExportFormat::Json => "json",
            ExportFormat::Html => "html",
            ExportFormat::Warc => "warc",
            ExportFormat::Zip => "zip",
            ExportFormat::Tar => "tar",
            ExportFormat::AgentTar => "agent-tar",
            ExportFormat::All => "all",
        }
    }
}

pub fn export_archive(
    config: &Config,
    archive_id: &str,
    format: ExportFormat,
    output: Option<&Path>,
) -> Result<PathBuf> {
    let layout = layout_for(config, archive_id)?;
    ensure_archive_exists(&layout)?;
    let output = output
        .map(Path::to_path_buf)
        .unwrap_or_else(|| layout.exports_dir().join(format.as_str()));
    fs::create_dir_all(&output)?;

    match format {
        ExportFormat::Markdown => copy_dir(&layout.readable_markdown_dir(), &output)?,
        ExportFormat::BasicMarkdown => copy_dir(&layout.readable_basic_markdown_dir(), &output)?,
        ExportFormat::Json => copy_dir(&layout.readable_json_dir(), &output)?,
        ExportFormat::Html => copy_dir(&layout.raw_pages_dir(), &output)?,
        ExportFormat::Jsonl => {
            fs::copy(layout.chunks_jsonl(), output.join("chunks.jsonl"))?;
        }
        ExportFormat::Warc => {
            write_warc_like(&layout, &output.join(format!("{archive_id}.warc.jsonl")))?;
        }
        ExportFormat::Zip => {
            zip_archive(&layout, &output.join(format!("{archive_id}.zip")))?;
        }
        ExportFormat::Tar => {
            tar_archive(&layout, &output.join(format!("{archive_id}.tar")))?;
        }
        ExportFormat::AgentTar => {
            agent_tar_archive(&layout, &output.join(format!("{archive_id}-agent.tar.gz")))?;
        }
        ExportFormat::All => {
            copy_dir(&layout.readable_markdown_dir(), &output.join("markdown"))?;
            copy_dir(
                &layout.readable_basic_markdown_dir(),
                &output.join("basic_markdown"),
            )?;
            copy_dir(&layout.readable_json_dir(), &output.join("json"))?;
            copy_dir(&layout.raw_pages_dir(), &output.join("raw_pages"))?;
            fs::copy(layout.chunks_jsonl(), output.join("chunks.jsonl"))?;
            write_warc_like(&layout, &output.join(format!("{archive_id}.warc.jsonl")))?;
            zip_archive(&layout, &output.join(format!("{archive_id}.zip")))?;
            tar_archive(&layout, &output.join(format!("{archive_id}.tar")))?;
            agent_tar_archive(&layout, &output.join(format!("{archive_id}-agent.tar.gz")))?;
        }
    }
    Ok(output)
}

pub fn create_ai_packs(
    config: &Config,
    archive_id: &str,
    max_tokens: Option<usize>,
) -> Result<Vec<PathBuf>> {
    let layout = layout_for(config, archive_id)?;
    ensure_archive_exists(&layout)?;
    fs::create_dir_all(layout.packs_dir())?;
    for entry in fs::read_dir(layout.packs_dir())? {
        let entry = entry?;
        if entry.file_name().to_string_lossy().starts_with("pack_") {
            let _ = fs::remove_file(entry.path());
        }
    }

    let storage = Storage::open(&layout.database())?;
    let chunks_by_page = load_chunks_by_page(&layout)?;
    let pages = storage
        .list_pages()?
        .into_iter()
        .filter(|page| page.markdown_path.is_some())
        .collect::<Vec<_>>();
    let token_budget = max_tokens.unwrap_or(config.ai_pack_token_budget).max(1_000);
    let mut packs = Vec::new();
    let mut current = String::new();
    let mut current_tokens = 0usize;
    let mut current_toc: Vec<(String, String)> = Vec::new();
    let mut pack_index = 1usize;

    for page in pages {
        let Some(markdown_rel) = page.markdown_path.as_deref() else {
            continue;
        };
        let markdown_path = layout.base().join(markdown_rel);
        let markdown = fs::read_to_string(&markdown_path)?;
        let page_tokens = estimate_tokens(&markdown);

        let units = if page_tokens > token_budget {
            chunks_by_page
                .get(&page.page_id)
                .map(|chunks| chunk_pack_units(&page.title, &page.source_url, chunks))
                .filter(|units| !units.is_empty())
                .unwrap_or_else(|| {
                    vec![PackUnit::new(
                        page.title.clone(),
                        page.source_url.clone(),
                        format!(
                            "\n\n# {}\n\nSource: {}\n\n{}",
                            page.title, page.source_url, markdown
                        ),
                    )]
                })
        } else {
            vec![PackUnit::new(
                page.title.clone(),
                page.source_url.clone(),
                format!(
                    "\n\n# {}\n\nSource: {}\n\n{}",
                    page.title, page.source_url, markdown
                ),
            )]
        };

        for unit in units {
            if current_tokens > 0 && current_tokens + unit.tokens > token_budget {
                let path = write_pack(&layout, pack_index, &current_toc, &current)?;
                packs.push(path);
                pack_index += 1;
                current.clear();
                current_tokens = 0;
                current_toc.clear();
            }
            current_toc.push((unit.title, unit.source_url));
            current.push_str(&unit.body);
            current_tokens += unit.tokens;
        }
    }

    if !current.is_empty() {
        packs.push(write_pack(&layout, pack_index, &current_toc, &current)?);
    }
    Ok(packs)
}

struct PackUnit {
    title: String,
    source_url: String,
    body: String,
    tokens: usize,
}

impl PackUnit {
    fn new(title: String, source_url: String, body: String) -> Self {
        let tokens = estimate_tokens(&body);
        Self {
            title,
            source_url,
            body,
            tokens,
        }
    }
}

fn load_chunks_by_page(layout: &ArchiveLayout) -> Result<HashMap<String, Vec<JsonlChunk>>> {
    let mut chunks_by_page: HashMap<String, Vec<JsonlChunk>> = HashMap::new();
    if !layout.chunks_jsonl().exists() {
        return Ok(chunks_by_page);
    }
    let raw = fs::read_to_string(layout.chunks_jsonl())?;
    for line in raw.lines().filter(|line| !line.trim().is_empty()) {
        let chunk: JsonlChunk = serde_json::from_str(line)?;
        chunks_by_page
            .entry(chunk.page_id.clone())
            .or_default()
            .push(chunk);
    }
    for chunks in chunks_by_page.values_mut() {
        chunks.sort_by_key(|chunk| chunk.order_index);
    }
    Ok(chunks_by_page)
}

fn chunk_pack_units(title: &str, source_url: &str, chunks: &[JsonlChunk]) -> Vec<PackUnit> {
    chunks
        .iter()
        .map(|chunk| {
            let heading = if chunk.heading_path.is_empty() {
                title.to_string()
            } else {
                chunk.heading_path.join(" > ")
            };
            let body = if chunk.content_type == "code" {
                let language = chunk.code_language.as_deref().unwrap_or("");
                format!(
                    "\n\n# {}\n\nSource: {}\n\nHeading: {}\n\n```{}\n{}\n```\n",
                    title,
                    source_url,
                    heading,
                    language,
                    chunk.text.trim_end()
                )
            } else {
                format!(
                    "\n\n# {}\n\nSource: {}\n\nHeading: {}\n\n{}\n",
                    title,
                    source_url,
                    heading,
                    chunk.text.trim()
                )
            };
            PackUnit::new(
                format!("{} / chunk {}", title, chunk.order_index),
                source_url.to_string(),
                body,
            )
        })
        .collect()
}

fn write_pack(
    layout: &ArchiveLayout,
    index: usize,
    toc: &[(String, String)],
    body: &str,
) -> Result<PathBuf> {
    let path = layout.packs_dir().join(format!("pack_{index:04}.md"));
    let mut content = String::new();
    content.push_str("# Siteforge AI Context Pack\n\n");
    content.push_str("This pack preserves source attribution and local asset references. OCR-derived content is labeled in page Markdown.\n\n");
    content.push_str("## Table of Contents\n\n");
    for (title, url) in toc {
        content.push_str(&format!("- {} - {}\n", title, url));
    }
    content.push_str(body);
    fs::write(&path, content)?;
    Ok(path)
}

fn write_warc_like(layout: &ArchiveLayout, path: &Path) -> Result<()> {
    #[derive(Serialize)]
    struct Record {
        record_type: &'static str,
        source_url: String,
        final_url: String,
        status: i64,
        content_hash: Option<String>,
        raw_path: Option<String>,
    }

    let storage = Storage::open(&layout.database())?;
    let mut file = File::create(path)?;
    for page in storage.list_pages()? {
        let record = Record {
            record_type: "response",
            source_url: page.source_url,
            final_url: page.final_url,
            status: page.status,
            content_hash: page.content_hash,
            raw_path: page.raw_path,
        };
        writeln!(file, "{}", serde_json::to_string(&record)?)?;
    }
    Ok(())
}

fn zip_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
    let file = File::create(path)?;
    let mut zip = zip::ZipWriter::new(file);
    let options = SimpleFileOptions::default().compression_method(zip::CompressionMethod::Deflated);
    for entry in WalkDir::new(layout.base())
        .into_iter()
        .filter_map(std::result::Result::ok)
    {
        if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
            continue;
        }
        let relative = entry
            .path()
            .strip_prefix(layout.base())
            .unwrap_or(entry.path())
            .to_string_lossy()
            .replace('\\', "/");
        zip.start_file(relative, options)?;
        let mut source = File::open(entry.path())?;
        let mut bytes = Vec::new();
        source.read_to_end(&mut bytes)?;
        zip.write_all(&bytes)?;
    }
    zip.finish()?;
    Ok(())
}

fn tar_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
    let file = File::create(path)?;
    let mut builder = tar::Builder::new(file);
    for entry in WalkDir::new(layout.base())
        .into_iter()
        .filter_map(std::result::Result::ok)
    {
        if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
            continue;
        }
        let relative = entry
            .path()
            .strip_prefix(layout.base())
            .unwrap_or(entry.path());
        builder.append_path_with_name(entry.path(), relative)?;
    }
    builder.finish()?;
    Ok(())
}

fn agent_tar_archive(layout: &ArchiveLayout, path: &Path) -> Result<()> {
    let file = File::create(path)?;
    let encoder = GzEncoder::new(file, Compression::default());
    let mut builder = tar::Builder::new(encoder);
    for entry in WalkDir::new(layout.base())
        .into_iter()
        .filter_map(std::result::Result::ok)
    {
        if !entry.file_type().is_file() || should_skip_bundle_path(layout, entry.path(), path) {
            continue;
        }
        let relative = entry
            .path()
            .strip_prefix(layout.base())
            .unwrap_or(entry.path());
        let relative_text = relative.to_string_lossy().replace('\\', "/");
        if !is_agent_bundle_file(&relative_text) {
            continue;
        }
        builder.append_path_with_name(entry.path(), relative)?;
    }
    builder.finish()?;
    let encoder = builder.into_inner()?;
    encoder.finish()?;
    Ok(())
}

fn should_skip_bundle_path(layout: &ArchiveLayout, entry_path: &Path, output_path: &Path) -> bool {
    if entry_path == output_path {
        return true;
    }
    let relative = entry_path
        .strip_prefix(layout.base())
        .unwrap_or(entry_path)
        .to_string_lossy()
        .replace('\\', "/");
    relative.starts_with("exports/")
        || relative.ends_with(".db-wal")
        || relative.ends_with(".db-shm")
}

fn is_agent_bundle_file(relative: &str) -> bool {
    relative == "AGENTS.md"
        || relative == "agent-index.json"
        || relative == "manifest.json"
        || relative == "checksums.json"
        || relative.starts_with("readable/")
        || relative.starts_with("chunks/")
        || relative.starts_with("packs/")
        || relative.starts_with("raw/pages/")
        || relative.starts_with("raw/rendered/")
        || relative.starts_with("raw/assets/")
}

fn copy_dir(source: &Path, destination: &Path) -> Result<()> {
    fs::create_dir_all(destination)?;
    if !source.exists() {
        return Ok(());
    }
    for entry in WalkDir::new(source)
        .into_iter()
        .filter_map(std::result::Result::ok)
    {
        let relative = entry.path().strip_prefix(source).unwrap_or(entry.path());
        let target = destination.join(relative);
        if entry.file_type().is_dir() {
            fs::create_dir_all(target)?;
        } else if entry.file_type().is_file() {
            if let Some(parent) = target.parent() {
                fs::create_dir_all(parent)?;
            }
            fs::copy(entry.path(), target)?;
        }
    }
    Ok(())
}

fn layout_for(config: &Config, archive_id: &str) -> Result<ArchiveLayout> {
    Ok(ArchiveLayout::new(
        config.resolved_archive_root()?,
        archive_id.to_string(),
    ))
}

fn ensure_archive_exists(layout: &ArchiveLayout) -> Result<()> {
    if !layout.base().exists() {
        return Err(SiteforgeError::ArchiveNotFound(layout.archive_id.clone()));
    }
    Ok(())
}
