import Foundation
import MarkBase

// ─────────────────────────────────────────────────────────────
// MarkBase CLI - Command Line Interface
// ─────────────────────────────────────────────────────────────

/// CLI command
public enum CLICommand {
    case serve(ServeOptions)
    case chat(ChatOptions)
    case list(ListOptions)
    case load(LoadOptions)
    case unload(UnloadOptions)
    case switchModel(SwitchOptions)
    case download(DownloadOptions)
    case search(SearchOptions)
    case benchmark(BenchmarkOptions)
    case help
}

/// CLI options
public struct ServeOptions {
    public var modelDir: String
    public var port: Int
    public var host: String
    public var maxConcurrency: Int
    
    public init(
        modelDir: String = "./model",
        port: Int = 8080,
        host: String = "127.0.0.1",
        maxConcurrency: Int = 4
    ) {
        self.modelDir = modelDir
        self.port = port
        self.host = host
        self.maxConcurrency = maxConcurrency
    }
}

public struct ChatOptions {
    public var modelDir: String
    public var prompt: String
    public var maxTokens: Int
    public var temperature: Float
    public var stream: Bool
    
    public init(
        modelDir: String = "./model",
        prompt: String = "",
        maxTokens: Int = 100,
        temperature: Float = 0.7,
        stream: Bool = true
    ) {
        self.modelDir = modelDir
        self.prompt = prompt
        self.maxTokens = maxTokens
        self.temperature = temperature
        self.stream = stream
    }
}

public struct ListOptions {
    public var modelsDir: String
    
    public init(modelsDir: String = "./models") {
        self.modelsDir = modelsDir
    }
}

public struct LoadOptions {
    public var modelId: String
    public var modelDir: String
    
    public init(modelId: String, modelDir: String = "./models") {
        self.modelId = modelId
        self.modelDir = modelDir
    }
}

public struct UnloadOptions {
    public var modelId: String?
    
    public init(modelId: String? = nil) {
        self.modelId = modelId
    }
}

public struct SwitchOptions {
    public var modelId: String
    
    public init(modelId: String) {
        self.modelId = modelId
    }
}

public struct DownloadOptions {
    public var repoId: String
    public var outputDir: String
    
    public init(repoId: String, outputDir: String = "./models") {
        self.repoId = repoId
        self.outputDir = outputDir
    }
}

public struct SearchOptions {
    public var query: String
    public var limit: Int
    public var gguf: Bool
    public var safetensors: Bool
    
    public init(
        query: String = "",
        limit: Int = 20,
        gguf: Bool = false,
        safetensors: Bool = false
    ) {
        self.query = query
        self.limit = limit
        self.gguf = gguf
        self.safetensors = safetensors
    }
}

public struct BenchmarkOptions {
    public var modelDir: String
    public var modelName: String
    public var numPrompts: Int
    
    public init(
        modelDir: String = "./model",
        modelName: String = "markbase",
        numPrompts: Int = 10
    ) {
        self.modelDir = modelDir
        self.modelName = modelName
        self.numPrompts = numPrompts
    }
}

/// CLI parser
public final class CLIParser {
    public static func parse(arguments: [String]) -> CLICommand {
        let args = Array(arguments.dropFirst()) // Skip program name
        
        guard !args.isEmpty else {
            return .help
        }
        
        let command = args[0]
        
        switch command {
        case "serve":
            return parseServe(args: Array(args.dropFirst()))
        case "chat":
            return parseChat(args: Array(args.dropFirst()))
        case "list":
            return .list(ListOptions())
        case "load":
            return parseLoad(args: Array(args.dropFirst()))
        case "unload":
            return .unload(UnloadOptions())
        case "switch":
            return parseSwitch(args: Array(args.dropFirst()))
        case "download":
            return parseDownload(args: Array(args.dropFirst()))
        case "search":
            return parseSearch(args: Array(args.dropFirst()))
        case "benchmark":
            return parseBenchmark(args: Array(args.dropFirst()))
        case "help", "--help", "-h":
            return .help
        default:
            print("Unknown command: \(command)")
            return .help
        }
    }
    
    private static func parseServe(args: [String]) -> CLICommand {
        var options = ServeOptions()
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--model", "-m":
                if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--port", "-p":
                if i + 1 < args.count, let port = Int(args[i + 1]) { options.port = port; i += 2 }
                else { i += 1 }
            case "--host":
                if i + 1 < args.count { options.host = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--concurrency", "-c":
                if i + 1 < args.count, let concurrency = Int(args[i + 1]) { options.maxConcurrency = concurrency; i += 2 }
                else { i += 1 }
            default:
                i += 1
            }
        }
        
        return .serve(options)
    }
    
    private static func parseChat(args: [String]) -> CLICommand {
        var options = ChatOptions()
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--model", "-m":
                if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--prompt", "-p":
                if i + 1 < args.count { options.prompt = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--max-tokens", "-n":
                if i + 1 < args.count, let n = Int(args[i + 1]) { options.maxTokens = n; i += 2 }
                else { i += 1 }
            case "--temperature", "-t":
                if i + 1 < args.count, let t = Float(args[i + 1]) { options.temperature = t; i += 2 }
                else { i += 1 }
            case "--stream", "-s":
                options.stream = true; i += 1
            case "--no-stream":
                options.stream = false; i += 1
            default:
                i += 1
            }
        }
        
        return .chat(options)
    }
    
    private static func parseLoad(args: [String]) -> CLICommand {
        var modelId = ""
        var modelDir = "./models"
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--model", "-m":
                if i + 1 < args.count { modelId = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--dir", "-d":
                if i + 1 < args.count { modelDir = args[i + 1]; i += 2 }
                else { i += 1 }
            default:
                if modelId.isEmpty { modelId = args[i] }
                i += 1
            }
        }
        
        return .load(LoadOptions(modelId: modelId, modelDir: modelDir))
    }
    
    private static func parseSwitch(args: [String]) -> CLICommand {
        var modelId = ""
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--model", "-m":
                if i + 1 < args.count { modelId = args[i + 1]; i += 2 }
                else { i += 1 }
            default:
                if modelId.isEmpty { modelId = args[i] }
                i += 1
            }
        }
        
        return .switchModel(SwitchOptions(modelId: modelId))
    }
    
    private static func parseDownload(args: [String]) -> CLICommand {
        var repoId = ""
        var outputDir = "./models"
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--output", "-o":
                if i + 1 < args.count { outputDir = args[i + 1]; i += 2 }
                else { i += 1 }
            default:
                if repoId.isEmpty { repoId = args[i] }
                i += 1
            }
        }
        
        return .download(DownloadOptions(repoId: repoId, outputDir: outputDir))
    }
    
    private static func parseSearch(args: [String]) -> CLICommand {
        var options = SearchOptions()
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--query", "-q":
                if i + 1 < args.count { options.query = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--limit", "-l":
                if i + 1 < args.count, let limit = Int(args[i + 1]) { options.limit = limit; i += 2 }
                else { i += 1 }
            case "--gguf":
                options.gguf = true; i += 1
            case "--safetensors":
                options.safetensors = true; i += 1
            default:
                if options.query.isEmpty { options.query = args[i] }
                i += 1
            }
        }
        
        return .search(options)
    }
    
    private static func parseBenchmark(args: [String]) -> CLICommand {
        var options = BenchmarkOptions()
        
        var i = 0
        while i < args.count {
            switch args[i] {
            case "--model", "-m":
                if i + 1 < args.count { options.modelDir = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--name":
                if i + 1 < args.count { options.modelName = args[i + 1]; i += 2 }
                else { i += 1 }
            case "--prompts", "-n":
                if i + 1 < args.count, let n = Int(args[i + 1]) { options.numPrompts = n; i += 2 }
                else { i += 1 }
            default:
                i += 1
            }
        }
        
        return .benchmark(options)
    }
}

/// CLI help message
public func printHelp() {
    print("""
    MarkBase CLI - Command Line Interface
    
    Usage: markbase <command> [options]
    
    Commands:
      serve         Start API server
      chat          Interactive chat
      list          List available models
      load          Load a model
      unload        Unload current model
      switch        Switch to a different model
      download      Download model from HuggingFace
      search        Search models on HuggingFace
      benchmark     Run performance benchmark
      help          Show this help message
    
    Examples:
      markbase serve --model ./model --port 8080
      markbase chat --model ./model --prompt "Hello!"
      markbase download mlx-community/gemma-4-e4b-it-4bit
      markbase search "gemma-4" --gguf
      markbase benchmark --model ./model --prompts 10
    
    Use "markbase <command> --help" for more information about a command.
    """)
}