use std::fs;
use std::io::Write;
use std::path::PathBuf;
use std::str::FromStr;

use comemo::Tracked;
use typst::diag::{At, FileError, FileResult, SourceResult, StrResult, bail};
use typst::engine::Engine;
use typst::foundations::{
    Array, Bytes, Content, Context, Datetime, Duration, IntoValue, NativeElement,
    NoneValue, Packed, Repr, Smart, StyleChain, Value, elem, func,
};
use typst::introspection::Locator;
use typst::layout::{Abs, BlockElem, Fragment, Margin, PageElem, Regions};
use typst::model::{Numbering, NumberingPattern};
use typst::syntax::{FileId, Source, Span};
use typst::text::{Font, FontBook, TextElem, TextSize};
use typst::utils::{LazyHash, singleton};
use typst::visualize::Color;
use typst::{Features, Library, LibraryExt, World};
use typst_kit::datetime::Time;
use typst_kit::files::{FileLoader, FileStore};
use typst_layout::layout_fragment;
use typst_syntax::package::PackageSpec;
use typst_syntax::{RootedPath, VirtualPath, VirtualRoot};
use unscanny::Scanner;

/// A world that provides access to the tests environment.
#[derive(Clone)]
pub struct TestWorld {
    main: Source,
    base: &'static TestBase,
}

impl TestWorld {
    /// Create a new world for a single test.
    ///
    /// This is cheap because the shared base for all test runs is lazily
    /// initialized just once.
    pub fn new(source: Source) -> Self {
        Self {
            main: source,
            base: singleton!(TestBase, TestBase::default()),
        }
    }
}

impl World for TestWorld {
    fn library(&self) -> &LazyHash<Library> {
        &self.base.library
    }

    fn book(&self) -> &LazyHash<FontBook> {
        &self.base.book
    }

    fn main(&self) -> FileId {
        self.main.id()
    }

    fn source(&self, id: FileId) -> FileResult<Source> {
        if id == self.main.id() {
            Ok(self.main.clone())
        } else {
            self.base.files.source(id)
        }
    }

    fn file(&self, id: FileId) -> FileResult<Bytes> {
        if id == self.main.id() {
            Ok(Bytes::from_string(self.main.clone()))
        } else {
            self.base.files.file(id)
        }
    }

    fn font(&self, index: usize) -> Option<Font> {
        self.base.fonts.get(index).cloned()
    }

    fn today(&self, offset: Option<Duration>) -> Option<Datetime> {
        let datetime = Datetime::from_ymd_hms(1970, 1, 1, 12, 0, 0).unwrap();
        Time::fixed(datetime).unwrap().today(offset)
    }
}

/// Shared foundation of all test worlds.
struct TestBase {
    library: LazyHash<Library>,
    book: LazyHash<FontBook>,
    fonts: Vec<Font>,
    files: FileStore<TestFiles>,
}

impl Default for TestBase {
    fn default() -> Self {
        let fonts: Vec<_> = typst_assets::fonts()
            .chain(typst_dev_assets::fonts())
            .flat_map(|data| Font::iter(Bytes::new(data)))
            .collect();

        Self {
            library: LazyHash::new(library()),
            book: LazyHash::new(FontBook::from_fonts(&fonts)),
            fonts,
            files: FileStore::new(TestFiles),
        }
    }
}

/// Loads files from the test suite, Typst assets, and the test packages.
/// Excludes the main source file, which is directly handled by the `World`.
pub struct TestFiles;

impl TestFiles {
    /// Resolves the file system path for a file ID.
    pub fn resolve(&self, id: FileId) -> PathBuf {
        let root = match id.root() {
            VirtualRoot::Project => PathBuf::new(),
            VirtualRoot::Package(spec) => {
                assert_eq!(spec.namespace, "test");
                format!("tests/packages/{}-{}", spec.name, spec.version).into()
            }
        };
        id.vpath().realize(&root)
    }

    /// Get the rooted path for a loaded file.
    pub fn rooted_path(path: &str) -> RootedPath {
        let mut s = Scanner::new(path);
        let root = if s.eat_if("tests/packages/") {
            let name = s.eat_until('-');
            s.expect('-');
            let version = s.eat_until('/');
            s.expect('/');
            VirtualRoot::Package(PackageSpec {
                namespace: "test".into(),
                name: name.into(),
                version: version.parse().unwrap(),
            })
        } else {
            VirtualRoot::Project
        };
        let vpath = VirtualPath::new(s.after()).unwrap();
        RootedPath::new(root, vpath)
    }
}

impl FileLoader for TestFiles {
    fn load(&self, id: FileId) -> FileResult<Bytes> {
        let path = self.resolve(id);

        // Resolve asset.
        if let Ok(suffix) = path.strip_prefix("assets/") {
            return typst_dev_assets::get(&suffix.to_string_lossy())
                .map(Bytes::new)
                .ok_or_else(|| FileError::NotFound(path));
        }

        let f = |e| FileError::from_io(e, &path);
        if fs::metadata(&path).map_err(f)?.is_dir() {
            Err(FileError::IsDirectory)
        } else {
            fs::read(&path).map(Bytes::new).map_err(f)
        }
    }
}

/// The extended standard library for testing.
fn library() -> Library {
    // Set page width to 120pt with 10pt margins, so that the inner page is
    // exactly 100pt wide. Page height is unbounded and font size is 10pt so
    // that it multiplies to nice round numbers.
    let mut lib = Library::builder().with_features(Features::all()).build();

    // Hook up helpers into the global scope.
    lib.global.scope_mut().define_func::<test>();
    lib.global.scope_mut().define_func::<test_repr>();
    lib.global.scope_mut().define_func::<print>();
    lib.global.scope_mut().define_func::<lines>();
    lib.global.scope_mut().define_func::<bounds>();
    lib.global
        .scope_mut()
        .define("conifer", Color::from_u8(0x9f, 0xEB, 0x52, 0xFF));
    lib.global
        .scope_mut()
        .define("forest", Color::from_u8(0x43, 0xA1, 0x27, 0xFF));

    // Hook up default styles.
    lib.styles.set(PageElem::width, Smart::Custom(Abs::pt(120.0).into()));
    lib.styles.set(PageElem::height, Smart::Auto);
    lib.styles
        .set(PageElem::margin, Margin::splat(Some(Smart::Custom(Abs::pt(10.0).into()))));
    lib.styles.set(TextElem::size, TextSize(Abs::pt(10.0).into()));

    lib
}

#[func]
fn test(lhs: Value, rhs: Value) -> StrResult<NoneValue> {
    if lhs != rhs {
        bail!("Assertion failed: {} != {}", lhs.repr(), rhs.repr());
    }
    Ok(NoneValue)
}

#[func]
fn test_repr(lhs: Value, rhs: Value) -> StrResult<NoneValue> {
    if lhs.repr() != rhs.repr() {
        bail!("Assertion failed: {} != {}", lhs.repr(), rhs.repr());
    }
    Ok(NoneValue)
}

#[func]
fn print(#[variadic] values: Vec<Value>) -> NoneValue {
    let mut out = std::io::stdout().lock();
    write!(out, "> ").unwrap();
    for (i, value) in values.into_iter().enumerate() {
        if i > 0 {
            write!(out, ", ").unwrap();
        }
        write!(out, "{value:?}").unwrap();
    }
    writeln!(out).unwrap();
    NoneValue
}

/// Generates `count` lines of text based on the numbering.
#[func]
fn lines(
    engine: &mut Engine,
    context: Tracked<Context>,
    span: Span,
    count: u64,
    #[default(Numbering::Pattern(NumberingPattern::from_str("A").unwrap()))]
    numbering: Numbering,
) -> SourceResult<Value> {
    (1..=count)
        .map(|n| numbering.apply(engine, context, span, &[n]))
        .collect::<SourceResult<Array>>()?
        .join(Some('\n'.into_value()), None, None)
        .at(span)
}

/// Display boundaries and the baseline around some content's frames.
#[func]
fn bounds(content: Content) -> SourceResult<Content> {
    Ok(BlockElem::multi_layouter(Packed::new(BoundsElem::new(content)), layout_bounds)
        .pack())
}

#[elem]
struct BoundsElem {
    #[positional]
    #[required]
    body: Content,
}

fn layout_bounds(
    elem: &Packed<BoundsElem>,
    engine: &mut Engine,
    locator: Locator,
    styles: StyleChain,
    regions: Regions,
) -> SourceResult<Fragment> {
    let mut fragment = layout_fragment(engine, &elem.body, locator, styles, regions)?;
    for frame in &mut fragment {
        frame.mark_box_in_place();
    }

    Ok(fragment)
}
