import Foundation

@Observable
final class AIRouter {
    private var clients: [AIProviderType: AIClient] = [:]
    private(set) var availability: [AIProviderType: Bool] = [:]

    // MARK: - Configuration

    func configure(with config: ProviderConfig) {
        clients.removeAll()

        for provider in AIProviderType.allCases where provider != .auto {
            let key = config.apiKey(for: provider)
            let model = config.model(for: provider)
            let endpoint = config.endpoint(for: provider)

            // Skip providers that need an API key but have none
            if provider.requiresAPIKey && key.isEmpty { continue }

            switch provider.clientKind {
            case .ollama:
                clients[provider] = OllamaClient(endpoint: endpoint, model: model)

            case .pollinations:
                clients[provider] = PollinationsClient(model: model)

            case .anthropic:
                clients[provider] = AnthropicClient(
                    apiKey: key, model: model, baseURL: endpoint
                )

            case .google:
                clients[provider] = GoogleGeminiClient(apiKey: key, model: model)

            case .openaiCompatible:
                clients[provider] = OpenAICompatibleClient(
                    provider: provider,
                    apiKey: key,
                    model: model,
                    baseURL: endpoint
                )
            }
        }
        print("[AIRouter] configure() — clients created: \(clients.keys.map(\.rawValue).sorted())")
    }

    // MARK: - Availability

    func refreshAvailability() async {
        await withTaskGroup(of: (AIProviderType, Bool).self) { group in
            for (type, client) in clients {
                group.addTask {
                    let ok = await client.checkAvailability()
                    return (type, ok)
                }
            }
            for await (type, ok) in group {
                availability[type] = ok
            }
        }
    }

    func isAvailable(_ provider: AIProviderType) -> Bool {
        if provider == .auto { return true }
        return availability[provider] ?? false
    }

    // MARK: - Routing

    /// Resolve `.auto` to the best concrete provider.
    /// Priority: Fast inference (free tiers) > Cloud (with keys) > Local > Aggregator > Free fallback
    func resolveProvider(_ selected: AIProviderType) -> AIProviderType {
        guard selected == .auto else { return selected }

        let priority: [AIProviderType] = [
            // Fast inference first — sub-second, usually free tier
            .groq, .cerebras, .sambanova,
            // Cloud providers with real models
            .openai, .anthropic, .google, .mistral, .xai, .deepseek,
            // Local (always available if running, but shouldn't override cloud if user set keys)
            .ollama, .lmstudio,
            // Other cloud
            .fireworks, .together, .lepton,
            // Aggregators
            .openrouter, .deepinfra, .perplexity, .huggingface, .replicate,
            .cohere, .ai21,
            // Free fallback
            .pollinations
        ]

        for p in priority {
            if availability[p] == true { return p }
        }
        return .pollinations
    }

    // MARK: - Completion

    func complete(
        prompt: String,
        systemPrompt: String,
        provider: AIProviderType
    ) async throws -> (String, AIProviderType) {
        let resolved = resolveProvider(provider)
        print("[AIRouter] complete() called — selected: \(provider.rawValue), resolved: \(resolved.rawValue), client exists: \(clients[resolved] != nil)")
        guard let client = clients[resolved] else {
            throw AIError.providerUnavailable
        }
        let response = try await client.complete(prompt: prompt, systemPrompt: systemPrompt)
        return (response, resolved)
    }

    // MARK: - Streaming

    func stream(
        prompt: String,
        systemPrompt: String,
        provider: AIProviderType
    ) -> (AsyncThrowingStream<String, Error>, AIProviderType) {
        let resolved = resolveProvider(provider)
        print("[AIRouter] stream() called — selected: \(provider.rawValue), resolved: \(resolved.rawValue), client exists: \(clients[resolved] != nil)")
        guard let client = clients[resolved] else {
            // Build a clear error message
            let reason: String
            if provider == .auto {
                reason = "No providers are configured. Go to Settings (⌘,) and add an API key."
            } else if provider.requiresAPIKey {
                reason = "\(provider.rawValue) requires an API key. Go to Settings (⌘,) → \(provider.rawValue) and enter your key."
            } else {
                reason = "\(provider.rawValue) is not available. Check that it's running."
            }
            let errStream = AsyncThrowingStream<String, Error> { continuation in
                continuation.yield("[Provider Error] \(reason)")
                continuation.finish()
            }
            return (errStream, resolved)
        }
        return (client.stream(prompt: prompt, systemPrompt: systemPrompt), resolved)
    }
}
