package re3lib;

import java.io.File;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;

import org.sqlite.JDBC;

import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;

public class FunctionDatabase implements AutoCloseable {
  public static enum Type {
    Auto(0),
    Fix(1),
    Stub(2),
    Ref(3);

    private final int value;

    Type(int value) {
      this.value = value;
    }

    public int getValue() {
      return value;
    }

    public static Type fromValue(int value) {
      for (Type type : Type.values()) {
        if (type.value == value) {
          return type;
        }
      }
      throw new IllegalArgumentException("Unknown type value: " + value);
    }
  }

  public static class FunctionEntry {
    public Address address;
    public String name;
    public File file;
    public Type type;

    public FunctionEntry(Address address, String name, File file, Type type) {
      this.address = address;
      this.name = name;
      this.file = file;
      this.type = type;
    }
  }

  public static class GlobalEntry {
    public Address address;
    public String name;
    public String dataType;
    public File file;

    public GlobalEntry(Address address, String name, String dataType, File file) {
      this.address = address;
      this.name = name;
      this.dataType = dataType;
      this.file = file;
    }
  }

  private File dbFile;
  private transient GhidraScript script;
  private Connection connection;

  // Prepared statements for better performance
  private PreparedStatement findByNameFunctions;
  private PreparedStatement findByNameImports;
  private PreparedStatement findByAddressFunctions;
  private PreparedStatement findByAddressImports;
  private PreparedStatement insertOrReplaceFunctions;
  private PreparedStatement deleteByFilepathFunctions;
  private PreparedStatement deleteByFilepathImports;
  private PreparedStatement loadAllFunctions;
  private PreparedStatement loadAllImports;

  // Add these prepared statements after the existing ones
  private PreparedStatement findByNameGlobals;
  private PreparedStatement findByAddressGlobals;
  private PreparedStatement insertOrReplaceGlobals;
  private PreparedStatement deleteByFilepathGlobals;
  private PreparedStatement loadAllGlobals;

  public FunctionDatabase(GhidraScript script) {
    this.script = script;
    dbFile = RemanConfig.INSTANCE.databasePath;
    try {
      java.sql.DriverManager.registerDriver(new JDBC());
    } catch (SQLException e) {
      script.printerr("Error registering JDBC driver: " + e.getMessage());
    }
  }

  public void connect() throws Exception {
    if (connection != null && !connection.isClosed()) {
      return; // Already connected
    }

    if (!dbFile.exists()) {
      script.println("Database file not found: " + dbFile);
      // Create parent directories if they don't exist
      dbFile.getParentFile().mkdirs();
    }

    try {
      connection = DriverManager.getConnection("jdbc:sqlite:" + dbFile.getAbsolutePath());
      createTablesIfNotExist();
      prepareCachedStatements();
      script.println("Connected to database: " + dbFile);
    } catch (SQLException e) {
      script.println("Error connecting to database: " + e.getMessage());
      throw new Exception("Failed to connect to database", e);
    }
  }

  public void disconnect() throws Exception {
    if (connection != null && !connection.isClosed()) {
      try {
        // Close prepared statements
        closePreparedStatements();
        connection.close();
        script.println("Disconnected from database");
      } catch (SQLException e) {
        script.println("Error disconnecting from database: " + e.getMessage());
        throw new Exception("Failed to disconnect from database", e);
      }
    }
  }

  private void ensureConnection() throws Exception {
    if (connection == null || connection.isClosed()) {
      connect();
    }
  }

  private void prepareCachedStatements() throws SQLException {
    // Find by name
    findByNameFunctions = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Functions WHERE name = ?");
    findByNameImports = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Imports WHERE name = ?");

    // Find by address
    findByAddressFunctions = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Functions WHERE address = ?");
    findByAddressImports = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Imports WHERE address = ?");

    // Insert or replace
    insertOrReplaceFunctions = connection.prepareStatement(
        "INSERT OR REPLACE INTO Functions (filepath, name, address, type) VALUES (?, ?, ?, ?)");

    // Delete by filepath
    deleteByFilepathFunctions = connection.prepareStatement(
        "DELETE FROM Functions WHERE filepath = ?");
    deleteByFilepathImports = connection.prepareStatement(
        "DELETE FROM Imports WHERE filepath = ?");

    // Load all entries
    loadAllFunctions = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Functions");
    loadAllImports = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Imports");

    // Global statements
    findByNameGlobals = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Globals WHERE name = ?");
    findByAddressGlobals = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Globals WHERE address = ?");
    insertOrReplaceGlobals = connection.prepareStatement(
        "INSERT OR REPLACE INTO Globals (filepath, name, address, type) VALUES (?, ?, ?, ?)");
    deleteByFilepathGlobals = connection.prepareStatement(
        "DELETE FROM Globals WHERE filepath = ?");
    loadAllGlobals = connection.prepareStatement(
        "SELECT filepath, name, address, type FROM Globals");
  }

  private void closePreparedStatements() throws SQLException {
    if (findByNameFunctions != null)
      findByNameFunctions.close();
    if (findByNameImports != null)
      findByNameImports.close();
    if (findByAddressFunctions != null)
      findByAddressFunctions.close();
    if (findByAddressImports != null)
      findByAddressImports.close();
    if (insertOrReplaceFunctions != null)
      insertOrReplaceFunctions.close();
    if (deleteByFilepathFunctions != null)
      deleteByFilepathFunctions.close();
    if (deleteByFilepathImports != null)
      deleteByFilepathImports.close();
    if (loadAllFunctions != null)
      loadAllFunctions.close();
    if (loadAllImports != null)
      loadAllImports.close();
    
    if (findByNameGlobals != null)
      findByNameGlobals.close();
    if (findByAddressGlobals != null)
      findByAddressGlobals.close();
    if (insertOrReplaceGlobals != null)
      insertOrReplaceGlobals.close();
    if (deleteByFilepathGlobals != null)
      deleteByFilepathGlobals.close();
    if (loadAllGlobals != null)
      loadAllGlobals.close();
  }

  public List<FunctionEntry> loadAllEntries() throws Exception {
    ensureConnection();
    List<FunctionEntry> entries = new ArrayList<>();

    try {
      // Load from Functions table
      try (ResultSet rs = loadAllFunctions.executeQuery()) {
        while (rs.next()) {
          FunctionEntry entry = createEntryFromResultSet(rs);
          if (entry != null) {
            entries.add(entry);
          }
        }
      }

      script.println("Loaded " + entries.size() + " function entries from database");
      return entries;
    } catch (SQLException e) {
      script.println("Error loading entries: " + e.getMessage());
      throw new Exception("Failed to load entries", e);
    }
  }

  private FunctionEntry createEntryFromResultSet(ResultSet rs) throws SQLException {
    String filepath = rs.getString("filepath");
    String name = rs.getString("name");
    String addressStr = rs.getString("address");
    int typeValue = rs.getInt("type");

    if (addressStr != null && !addressStr.isEmpty()) {
      Address address = script.getCurrentProgram().getAddressFactory().getAddress(addressStr);
      File file = new File(RemanConfig.INSTANCE.outputDir, filepath);
      Type type = Type.fromValue(typeValue);

      return new FunctionEntry(address, name, file, type);
    }
    return null;
  }

  private void createTablesIfNotExist() throws SQLException {
    String createFunctions = """
        CREATE TABLE IF NOT EXISTS Functions (
            filepath TEXT,
            name TEXT,
            address TEXT,
            type INTEGER,
            PRIMARY KEY (name, filepath)
        )""";

    String createImports = """
        CREATE TABLE IF NOT EXISTS Imports (
            filepath TEXT,
            name TEXT,
            address TEXT,
            type INTEGER,
            PRIMARY KEY (name, filepath)
        )""";

    String createGlobals = """
        CREATE TABLE IF NOT EXISTS Globals (
            filepath TEXT,
            name TEXT,
            address TEXT,
            type TEXT,
            PRIMARY KEY (name, filepath)
        )""";

    connection.prepareStatement(createFunctions).executeUpdate();
    connection.prepareStatement(createImports).executeUpdate();
    connection.prepareStatement(createGlobals).executeUpdate();
  }

  // Helper method to find entries by name
  public List<FunctionEntry> findEntriesByName(String name) throws Exception {
    ensureConnection();
    List<FunctionEntry> results = new ArrayList<>();

    try {
      // Search Functions table
      findByNameFunctions.setString(1, name);
      try (ResultSet rs = findByNameFunctions.executeQuery()) {
        while (rs.next()) {
          FunctionEntry entry = createEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      // Search Imports table
      findByNameImports.setString(1, name);
      try (ResultSet rs = findByNameImports.executeQuery()) {
        while (rs.next()) {
          FunctionEntry entry = createEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      return results;
    } catch (SQLException e) {
      script.println("Error finding entries by name: " + e.getMessage());
      throw new Exception("Failed to find entries by name", e);
    }
  }

  // Helper method to find entries by address
  public List<FunctionEntry> findEntriesByAddress(Address address) throws Exception {
    ensureConnection();
    List<FunctionEntry> results = new ArrayList<>();
    String addressStr = address.toString();

    try {
      // Search Functions table
      findByAddressFunctions.setString(1, addressStr);
      try (ResultSet rs = findByAddressFunctions.executeQuery()) {
        while (rs.next()) {
          FunctionEntry entry = createEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      // Search Imports table
      findByAddressImports.setString(1, addressStr);
      try (ResultSet rs = findByAddressImports.executeQuery()) {
        while (rs.next()) {
          FunctionEntry entry = createEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      return results;
    } catch (SQLException e) {
      script.println("Error finding entries by address: " + e.getMessage());
      throw new Exception("Failed to find entries by address", e);
    }
  }

  // Helper method to add/update entry (insert or replace based on filename)
  public void addEntryAt(FunctionEntry entry) throws Exception {
    ensureConnection();

    String relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath()
        .relativize(entry.file.toPath()).toString().replace('\\', '/');

    try {
      insertOrReplaceFunctions.setString(1, relativePath);
      insertOrReplaceFunctions.setString(2, entry.name);
      insertOrReplaceFunctions.setString(3, entry.address.toString());
      insertOrReplaceFunctions.setInt(4, entry.type.getValue());
      insertOrReplaceFunctions.executeUpdate();

      script.println("Added/updated entry: " + entry.name + " at " + entry.address + " in " + relativePath);
    } catch (SQLException e) {
      script.println("Error adding entry: " + e.getMessage());
      throw new Exception("Failed to add entry", e);
    }
  }

  // Helper method to remove entry by file path
  public void removeEntryAt(String filePath) throws Exception {
    ensureConnection();

    String relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath()
        .relativize(new File(filePath).toPath()).toString().replace('\\', '/');

    try {
      deleteByFilepathFunctions.setString(1, relativePath);
      int deletedCount = deleteByFilepathFunctions.executeUpdate();

      deleteByFilepathImports.setString(1, relativePath);
      deletedCount += deleteByFilepathImports.executeUpdate();

      script.println("Removed " + deletedCount + " entries for file: " + relativePath);
    } catch (SQLException e) {
      script.println("Error removing entries: " + e.getMessage());
      throw new Exception("Failed to remove entries", e);
    }
  }

  public void add(FunctionEntry entry) throws Exception {
    // Add entry directly to database
    addEntryAt(entry);
  }

  public void applyDefaultFilters(boolean rebuildAllGlobals) throws Exception {
    GlobalDumper globalDumper = new GlobalDumper(script, this);
    FunctionDumper dumper = new FunctionDumper(script, this, globalDumper);

    if (rebuildAllGlobals) {
      globalDumper.removeGlobalManifest();
    }

    boolean madeAnyChanges = false;

    // Load all entries from database
    List<FunctionEntry> entries = loadAllEntries();

    // Create a hash map to store symbol names
    Map<Address, String> symbolNames = new HashMap<>();
    Map<String, File> exportedFunctionNames = new HashMap<>();
    for (FunctionEntry entry : entries) {
      Function function = script.getFunctionAt(entry.address);
      if (function != null) {
        boolean isAuto = entry.type == Type.Auto;
        boolean isFix = entry.type == Type.Fix;
        // Get the actual symbol name and store it in the hash map
        String symbolName = function.getName();
        symbolNames.put(entry.address, symbolName);

        if (isAuto && !exportedFunctionNames.containsKey(entry.name)) {
          exportedFunctionNames.put(entry.name, entry.file);
        } else if (isFix) {
          exportedFunctionNames.replace(entry.name, entry.file);
        }
      }
    }

    // Print the number of symbol names collected
    script.println("Collected " + symbolNames.size() + " symbol names");

    boolean dryMode = false;

    HashSet<Function> functionsToRegenerate = new HashSet<>();

    Iterator<FunctionEntry> iterator = entries.iterator();
    while (iterator.hasNext()) {
      FunctionEntry entry = iterator.next();
      Function function = script.getFunctionAt(entry.address);

      boolean pendingDelete = false;
      boolean pendingRegenerate = false;

      if (rebuildAllGlobals) {
        pendingRegenerate = true;
      }

      // Remove CRT and other blacklisted functions
      if (function == null || !dumper.isValidFunction(function)) {
        // Remove the file
        if (entry.file != null && entry.file.exists()) {
          script.println("Removed file: " + entry.file.getAbsolutePath());
          pendingDelete = true;
        }

        // Remove entry from the list
        script.println("Removed invalid function entry: " + entry.name + " at " + entry.address);
        function = null;
      }

      // Check if symbol name matches the symbol name parsed from the file
      if (function != null) {
        String actualSymbolName = symbolNames.get(entry.address);
        if (actualSymbolName == null) {
          throw new Exception(
              "Symbol name not found for function at " + entry.address + " in file " + entry.file.getAbsolutePath());
        }

        if (actualSymbolName != null && !actualSymbolName.equals(entry.name)) {
          File fnExportedFile = exportedFunctionNames.get(entry.name);
          if (fnExportedFile != null && fnExportedFile != entry.file) {
            // Already exists elsewhere, so remove this file
            script.println("Removing duplicate function: " + entry.name + " at " + entry.address + " overridden by "
                + fnExportedFile);
            pendingDelete = true;
          } else {
            // Regeneral this function
            script.println("Symbol name mismatch for function at " + entry.address + ": " +
                "File name: " + entry.name + ", Actual symbol: " + actualSymbolName);
            entry.name = actualSymbolName; // Update the entry name to match the actual symbol
            pendingRegenerate = true;
          }
        }

        entry.name = actualSymbolName; // Update the entry name to match the actual symbol
        madeAnyChanges = true;
      }

      if (pendingDelete) {
        iterator.remove();
        if (!dryMode) {
          entry.file.delete();
          // Remove from database
          removeEntryAt(entry.file.getAbsolutePath());
          madeAnyChanges = true;
        }
      } else if (pendingRegenerate && entry.type != Type.Stub) {
        if (!dryMode) {
          functionsToRegenerate.add(function);
          // Update entry in database with corrected name
          addEntryAt(entry);
          madeAnyChanges = true;
        }
      }
    }

    for (Function function : functionsToRegenerate) {
      script.println("Regenerating function: " + function.getName() + " at " + function.getEntryPoint());
      dumper.dump(function);
    }

    if (madeAnyChanges) {
      // Update CMake timestamp
      RemanConfig.INSTANCE.touchCMakeTimestamp();
      globalDumper.dumpGlobals();
      globalDumper.saveGlobalManifest();

      TypeDumper typeDumper = new TypeDumper(script);
      typeDumper.run();
    }
  }

  // Global-specific methods
  public List<GlobalEntry> loadAllGlobals() throws Exception {
    ensureConnection();
    List<GlobalEntry> globals = new ArrayList<>();

    try {
      try (ResultSet rs = loadAllGlobals.executeQuery()) {
        while (rs.next()) {
          GlobalEntry entry = createGlobalEntryFromResultSet(rs);
          if (entry != null) {
            globals.add(entry);
          }
        }
      }

      script.println("Loaded " + globals.size() + " global entries from database");
      return globals;
    } catch (SQLException e) {
      script.println("Error loading globals: " + e.getMessage());
      throw new Exception("Failed to load globals", e);
    }
  }

  private GlobalEntry createGlobalEntryFromResultSet(ResultSet rs) throws SQLException {
    String filepath = rs.getString("filepath");
    String name = rs.getString("name");
    String addressStr = rs.getString("address");
    String typeStr = rs.getString("type");

    if (addressStr != null && !addressStr.isEmpty()) {
      Address address = script.getCurrentProgram().getAddressFactory().getAddress(addressStr);
      File file = new File(RemanConfig.INSTANCE.outputDir, filepath);
      
      return new GlobalEntry(address, name, typeStr, file);
    }
    return null;
  }

  public List<GlobalEntry> findGlobalsByName(String name) throws Exception {
    ensureConnection();
    List<GlobalEntry> results = new ArrayList<>();

    try {
      findByNameGlobals.setString(1, name);
      try (ResultSet rs = findByNameGlobals.executeQuery()) {
        while (rs.next()) {
          GlobalEntry entry = createGlobalEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      return results;
    } catch (SQLException e) {
      script.println("Error finding globals by name: " + e.getMessage());
      throw new Exception("Failed to find globals by name", e);
    }
  }

  public List<GlobalEntry> findGlobalsByAddress(Address address) throws Exception {
    ensureConnection();
    List<GlobalEntry> results = new ArrayList<>();
    String addressStr = address.toString();

    try {
      findByAddressGlobals.setString(1, addressStr);
      try (ResultSet rs = findByAddressGlobals.executeQuery()) {
        while (rs.next()) {
          GlobalEntry entry = createGlobalEntryFromResultSet(rs);
          if (entry != null) {
            results.add(entry);
          }
        }
      }

      return results;
    } catch (SQLException e) {
      script.println("Error finding globals by address: " + e.getMessage());
      throw new Exception("Failed to find globals by address", e);
    }
  }

  public void addGlobal(Address address, String name, String dataType) throws Exception {
    ensureConnection();

    String filepath = RemanConfig.GLOBAL_H_FILE; // Default filepath for globals
    String addressStr = address.toString();

    try {
      insertOrReplaceGlobals.setString(1, filepath);
      insertOrReplaceGlobals.setString(2, name);
      insertOrReplaceGlobals.setString(3, addressStr);
      insertOrReplaceGlobals.setString(4, dataType);
      insertOrReplaceGlobals.executeUpdate();

      script.println("Added/updated global: " + name + " at " + address + " with type " + dataType);
    } catch (SQLException e) {
      script.println("Error adding global: " + e.getMessage());
      throw new Exception("Failed to add global", e);
    }
  }

  public void removeGlobalsByFilepath(String filePath) throws Exception {
    ensureConnection();

    String relativePath;
    
    // Check if filePath is already a relative path or just a filename
    File inputFile = new File(filePath);
    if (inputFile.isAbsolute()) {
      // Convert absolute path to relative
      try {
        relativePath = new File(RemanConfig.INSTANCE.outputDir).toPath()
            .relativize(inputFile.toPath()).toString().replace('\\', '/');
      } catch (IllegalArgumentException e) {
        // Fallback if paths can't be relativized (different drives, etc.)
        script.println("Warning: Could not relativize path: " + filePath + ", using as-is");
        relativePath = filePath.replace('\\', '/');
      }
    } else {
      // Already relative or just a filename, use as-is
      relativePath = filePath.replace('\\', '/');
    }

    try {
      deleteByFilepathGlobals.setString(1, relativePath);
      int deletedCount = deleteByFilepathGlobals.executeUpdate();

      script.println("Removed " + deletedCount + " global entries for file: " + relativePath);
    } catch (SQLException e) {
      script.println("Error removing global entries: " + e.getMessage());
      throw new Exception("Failed to remove global entries", e);
    }
  }

  @Override
  public void close() throws Exception {
    this.disconnect();
  }
}