import Foundation
import ImageIO
import Vision
import FoundationModels
import CoreVideo

struct Entrada: Decodable {
    let frame: String?
    let frames: [String]?
    let recursos: [String]
    let resumir: Bool
}

struct Evidencia: Encodable {
    let tipo: String
    let valor: [String: JSONValue]
    let confianca: Double?
    let modelo: String?
}

struct Saida: Encodable {
    let evidencias: [Evidencia]
    let avisos: [String]
}

enum JSONValue: Encodable {
    case texto(String), numero(Double), booleano(Bool), objeto([String: JSONValue]), lista([JSONValue]), nulo

    func encode(to encoder: Encoder) throws {
        var container = encoder.singleValueContainer()
        switch self {
        case .texto(let value): try container.encode(value)
        case .numero(let value): try container.encode(value)
        case .booleano(let value): try container.encode(value)
        case .objeto(let value): try container.encode(value)
        case .lista(let value): try container.encode(value)
        case .nulo: try container.encodeNil()
        }
    }
}

func caixa(_ rect: CGRect) -> [String: JSONValue] {
    ["x": .numero(rect.origin.x), "y": .numero(rect.origin.y),
     "largura": .numero(rect.size.width), "altura": .numero(rect.size.height)]
}

func ponto(_ point: CGPoint) -> JSONValue {
    .objeto(["x": .numero(point.x), "y": .numero(point.y)])
}

func ponto(_ point: VNRecognizedPoint) -> JSONValue {
    .objeto(["x": .numero(point.location.x), "y": .numero(point.location.y),
             "confianca": .numero(Double(point.confidence))])
}

func executar<T: VNRequest>(_ request: T, em url: URL) throws -> [VNObservation] {
    guard let source = CGImageSourceCreateWithURL(url as CFURL, nil),
          let imagem = CGImageSourceCreateImageAtIndex(source, 0, nil) else {
        throw NSError(domain: "JhonnyAppleVision", code: 1,
                      userInfo: [NSLocalizedDescriptionKey: "ImageIO não conseguiu decodificar o frame"])
    }
    let handler = VNImageRequestHandler(cgImage: imagem, options: [:])
    try handler.perform([request])
    return request.results ?? []
}

func coberturaDaMascara(_ pixelBuffer: CVPixelBuffer) -> Double {
    CVPixelBufferLockBaseAddress(pixelBuffer, .readOnly)
    defer { CVPixelBufferUnlockBaseAddress(pixelBuffer, .readOnly) }
    guard let base = CVPixelBufferGetBaseAddress(pixelBuffer) else { return 0 }
    let altura = CVPixelBufferGetHeight(pixelBuffer)
    let largura = CVPixelBufferGetWidth(pixelBuffer)
    let stride = CVPixelBufferGetBytesPerRow(pixelBuffer)
    let bytes = base.assumingMemoryBound(to: UInt8.self)
    var ocupados = 0
    for y in 0..<altura {
        for x in 0..<largura where bytes[y * stride + x] > 12 { ocupados += 1 }
    }
    return Double(ocupados) / Double(max(1, altura * largura))
}

func analisarVision(_ entrada: Entrada) -> Saida {
    guard let frame = entrada.frame else {
        return Saida(evidencias: [], avisos: ["frame: caminho obrigatório para análise individual"])
    }
    let url = URL(fileURLWithPath: frame)
    var evidencias: [Evidencia] = []
    var avisos: [String] = []

    func tentar(_ recurso: String, _ bloco: () throws -> [Evidencia]) {
        guard entrada.recursos.contains(recurso) else { return }
        do { evidencias.append(contentsOf: try bloco()) }
        catch { avisos.append("\(recurso): \(error.localizedDescription)") }
    }

    tentar("faces") {
        let request = VNDetectFaceLandmarksRequest()
        return try executar(request, em: url).compactMap { observacao -> Evidencia? in
            guard let face = observacao as? VNFaceObservation else { return nil }
            var valor: [String: JSONValue] = ["bounding_box": .objeto(caixa(face.boundingBox))]
            if let landmarks = face.landmarks {
                let grupos: [(String, VNFaceLandmarkRegion2D?)] = [
                    ("olho_esquerdo", landmarks.leftEye), ("olho_direito", landmarks.rightEye),
                    ("sobrancelha_esquerda", landmarks.leftEyebrow), ("sobrancelha_direita", landmarks.rightEyebrow),
                    ("nariz", landmarks.nose), ("labios", landmarks.outerLips), ("rosto", landmarks.faceContour)
                ]
                valor["landmarks"] = .objeto(Dictionary(uniqueKeysWithValues: grupos.compactMap { nome, regiao in
                    guard let regiao else { return nil }
                    return (nome, .lista(regiao.normalizedPoints.map { .objeto(["x": .numero($0.x), "y": .numero($0.y)]) }))
                }))
            }
            return Evidencia(tipo: "rosto", valor: valor, confianca: Double(face.confidence), modelo: "VNDetectFaceLandmarksRequest")
        }
    }

    tentar("qualidade_facial") {
        let request = VNDetectFaceCaptureQualityRequest()
        return try executar(request, em: url).compactMap { observacao -> Evidencia? in
            guard let face = observacao as? VNFaceObservation, let qualidade = face.faceCaptureQuality else { return nil }
            return Evidencia(tipo: "qualidade_do_rosto", valor: ["bounding_box": .objeto(caixa(face.boundingBox)),
                "score": .numero(Double(qualidade))], confianca: Double(face.confidence), modelo: "VNDetectFaceCaptureQualityRequest")
        }
    }

    tentar("pessoas") {
        let request = VNDetectHumanRectanglesRequest()
        request.upperBodyOnly = false
        return try executar(request, em: url).compactMap { observacao in
            guard let pessoa = observacao as? VNHumanObservation else { return nil }
            return Evidencia(tipo: "pessoa", valor: ["bounding_box": .objeto(caixa(pessoa.boundingBox)),
                "ocupacao_do_quadro": .numero(pessoa.boundingBox.width * pessoa.boundingBox.height)],
                confianca: Double(pessoa.confidence), modelo: "VNDetectHumanRectanglesRequest")
        }
    }

    tentar("pose") {
        let request = VNDetectHumanBodyPoseRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let pose = observacao as? VNHumanBodyPoseObservation else { return nil }
            let pontos = try? pose.recognizedPoints(.all)
            let dados: [String: JSONValue] = pontos.map { Dictionary(uniqueKeysWithValues: $0.map { (String(describing: $0.key), ponto($0.value)) }) } ?? [:]
            return Evidencia(tipo: "pose", valor: ["pontos": .objeto(dados)], confianca: Double(pose.confidence), modelo: "VNDetectHumanBodyPoseRequest")
        }
    }

    tentar("maos") {
        let request = VNDetectHumanHandPoseRequest()
        request.maximumHandCount = 4
        return try executar(request, em: url).compactMap { observacao in
            guard let mao = observacao as? VNHumanHandPoseObservation else { return nil }
            let pontos = try? mao.recognizedPoints(.all)
            let dados: [String: JSONValue] = pontos.map { Dictionary(uniqueKeysWithValues: $0.map { (String(describing: $0.key), ponto($0.value)) }) } ?? [:]
            return Evidencia(tipo: "mao", valor: ["pontos": .objeto(dados)], confianca: Double(mao.confidence), modelo: "VNDetectHumanHandPoseRequest")
        }
    }

    tentar("ocr") {
        let request = VNRecognizeTextRequest()
        request.recognitionLevel = .accurate
        request.recognitionLanguages = ["pt-BR", "en-US"]
        return try executar(request, em: url).compactMap { observacao in
            guard let texto = observacao as? VNRecognizedTextObservation,
                  let candidato = texto.topCandidates(1).first else { return nil }
            return Evidencia(tipo: "ocr", valor: ["texto": .texto(candidato.string),
                "bounding_box": .objeto(caixa(texto.boundingBox))], confianca: Double(candidato.confidence), modelo: "VNRecognizeTextRequest")
        }
    }

    tentar("categorias") {
        let request = VNClassifyImageRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let categoria = observacao as? VNClassificationObservation, categoria.confidence >= 0.2 else { return nil }
            return Evidencia(tipo: "categoria", valor: ["identificador": .texto(categoria.identifier)],
                confianca: Double(categoria.confidence), modelo: "VNClassifyImageRequest")
        }
    }

    tentar("estetica") {
        let request = VNCalculateImageAestheticsScoresRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let score = observacao as? VNImageAestheticsScoresObservation else { return nil }
            return Evidencia(tipo: "estetica", valor: ["score_global": .numero(Double(score.overallScore))],
                confianca: nil, modelo: "VNCalculateImageAestheticsScoresRequest")
        }
    }

    tentar("codigos") {
        let request = VNDetectBarcodesRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let codigo = observacao as? VNBarcodeObservation else { return nil }
            return Evidencia(tipo: "codigo", valor: ["payload": .texto(codigo.payloadStringValue ?? ""),
                "simbologia": .texto(codigo.symbology.rawValue), "bounding_box": .objeto(caixa(codigo.boundingBox))],
                confianca: Double(codigo.confidence), modelo: "VNDetectBarcodesRequest")
        }
    }

    tentar("horizonte") {
        let request = VNDetectHorizonRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let horizonte = observacao as? VNHorizonObservation else { return nil }
            return Evidencia(tipo: "horizonte", valor: ["angulo_radianos": .numero(Double(horizonte.angle))],
                confianca: Double(horizonte.confidence), modelo: "VNDetectHorizonRequest")
        }
    }

    tentar("retangulos") {
        let request = VNDetectRectanglesRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let retangulo = observacao as? VNRectangleObservation else { return nil }
            return Evidencia(tipo: "retangulo", valor: ["bounding_box": .objeto(caixa(retangulo.boundingBox)),
                "cantos": .objeto(["superior_esquerdo": ponto(retangulo.topLeft),
                                     "superior_direito": ponto(retangulo.topRight),
                                     "inferior_esquerdo": ponto(retangulo.bottomLeft),
                                     "inferior_direito": ponto(retangulo.bottomRight)])],
                confianca: Double(retangulo.confidence), modelo: "VNDetectRectanglesRequest")
        }
    }

    tentar("contornos") {
        let request = VNDetectContoursRequest()
        return try executar(request, em: url).compactMap { observacao in
            guard let contornos = observacao as? VNContoursObservation else { return nil }
            return Evidencia(tipo: "contornos", valor: ["quantidade_principal": .numero(Double(contornos.topLevelContours.count))],
                confianca: Double(contornos.confidence), modelo: "VNDetectContoursRequest")
        }
    }

    tentar("segmentacao_de_pessoas") {
        let request = VNGeneratePersonSegmentationRequest()
        request.qualityLevel = .balanced
        return try executar(request, em: url).compactMap { observacao in
            guard let mascara = observacao as? VNPixelBufferObservation else { return nil }
            return Evidencia(tipo: "segmentacao_de_pessoas", valor: ["ocupacao_do_quadro": .numero(coberturaDaMascara(mascara.pixelBuffer))],
                confianca: nil, modelo: "VNGeneratePersonSegmentationRequest")
        }
    }

    return Saida(evidencias: evidencias, avisos: avisos)
}

func featurePrint(_ url: URL) throws -> VNFeaturePrintObservation {
    let request = VNGenerateImageFeaturePrintRequest()
    guard let resultado = try executar(request, em: url).first as? VNFeaturePrintObservation else {
        throw NSError(domain: "JhonnyAppleVision", code: 2, userInfo: [NSLocalizedDescriptionKey: "Vision não produziu feature print"])
    }
    return resultado
}

func analisarSequencia(_ entrada: Entrada) -> Saida {
    let caminhos = entrada.frames ?? []
    guard caminhos.count >= 2 else { return Saida(evidencias: [], avisos: ["sequencia: informe ao menos dois frames"]) }
    var evidencias: [Evidencia] = []
    var avisos: [String] = []
    for indice in 0..<(caminhos.count - 1) {
        do {
            var distancia: Float = 0
            try featurePrint(URL(fileURLWithPath: caminhos[indice])).computeDistance(
                &distancia, to: featurePrint(URL(fileURLWithPath: caminhos[indice + 1])))
            evidencias.append(Evidencia(tipo: "similaridade_visual", valor: [
                "frame_inicial": .numero(Double(indice)), "frame_final": .numero(Double(indice + 1)),
                "distancia_feature_print": .numero(Double(distancia)),
                "possivel_mudanca_de_cena": .booleano(distancia > 12),
            ], confianca: nil, modelo: "VNGenerateImageFeaturePrintRequest"))
        } catch { avisos.append("similaridade_visual[\(indice)]: \(error.localizedDescription)") }
    }
    return Saida(evidencias: evidencias, avisos: avisos)
}

extension Dictionary where Key == String, Value == JSONValue {
    func numero(_ chave: String) -> Double? {
        if case .numero(let valor)? = self[chave] { return valor }
        return nil
    }
    func texto(_ chave: String) -> String? {
        if case .texto(let valor)? = self[chave] { return valor }
        return nil
    }
    func booleano(_ chave: String) -> Bool? {
        if case .booleano(let valor)? = self[chave] { return valor }
        return nil
    }
    func objeto(_ chave: String) -> [String: JSONValue]? {
        if case .objeto(let valor)? = self[chave] { return valor }
        return nil
    }
}

/// Traduz uma evidência em uma frase factual com os valores medidos, nunca só o nome do tipo —
/// é isso que alimenta a interpretação editorial, então precisa carregar o fato, não só o rótulo.
func descreverEvidencia(_ evidencia: Evidencia) -> String {
    let valor = evidencia.valor
    switch evidencia.tipo {
    case "rosto":
        let grupos = valor.objeto("landmarks")?.count ?? 0
        return "rosto detectado (\(grupos) grupos de landmarks mapeados)"
    case "qualidade_do_rosto":
        guard let score = valor.numero("score") else { return "qualidade do rosto avaliada" }
        return "qualidade do rosto: score \(String(format: "%.2f", score))"
    case "pessoa":
        guard let ocupacao = valor.numero("ocupacao_do_quadro") else { return "pessoa detectada" }
        return "pessoa ocupando \(String(format: "%.0f", ocupacao * 100))% do quadro"
    case "pose":
        let pontos = valor.objeto("pontos")?.count ?? 0
        return "pose corporal com \(pontos) pontos reconhecidos"
    case "mao":
        let pontos = valor.objeto("pontos")?.count ?? 0
        return "mão com \(pontos) pontos reconhecidos"
    case "ocr":
        guard let texto = valor.texto("texto"), !texto.isEmpty else { return "nenhum texto legível na imagem" }
        return "texto detectado na imagem: \"\(texto)\""
    case "categoria":
        guard let identificador = valor.texto("identificador") else { return "categoria detectada" }
        let confianca = evidencia.confianca.map { String(format: "%.0f%%", $0 * 100) } ?? "confiança desconhecida"
        return "categoria \"\(identificador)\" (\(confianca))"
    case "estetica":
        guard let score = valor.numero("score_global") else { return "score estético calculado" }
        return "score estético global: \(String(format: "%.2f", score))"
    case "codigo":
        guard let payload = valor.texto("payload"), !payload.isEmpty else { return "código detectado sem payload legível" }
        return "código detectado: \"\(payload)\""
    case "horizonte":
        guard let angulo = valor.numero("angulo_radianos") else { return "horizonte detectado" }
        return "horizonte inclinado \(String(format: "%.1f", angulo * 180 / .pi))°"
    case "retangulo":
        guard let bbox = valor.objeto("bounding_box"), let largura = bbox.numero("largura"), let altura = bbox.numero("altura") else {
            return "retângulo/documento detectado no quadro"
        }
        return "retângulo detectado ocupando \(String(format: "%.0f", largura * 100))%x\(String(format: "%.0f", altura * 100))% do quadro"
    case "contornos":
        guard let quantidade = valor.numero("quantidade_principal") else { return "contornos detectados" }
        return "\(Int(quantidade)) contornos principais detectados"
    case "segmentacao_de_pessoas":
        guard let ocupacao = valor.numero("ocupacao_do_quadro") else { return "segmentação de pessoa calculada" }
        return "silhueta de pessoa cobrindo \(String(format: "%.0f", ocupacao * 100))% do quadro"
    case "similaridade_visual":
        let mudanca = valor.booleano("possivel_mudanca_de_cena") ?? false
        return mudanca ? "possível corte de cena entre os quadros" : "quadros visualmente semelhantes, sem corte de cena"
    default:
        return evidencia.tipo
    }
}

func interpretar(_ saida: Saida) async -> String? {
    guard SystemLanguageModel.default.isAvailable else { return nil }
    let fatos = saida.evidencias.map(descreverEvidencia).joined(separator: "; ")
    guard !fatos.isEmpty else { return nil }
    do {
        let sessao = LanguageModelSession(instructions: "Você é um assistente editorial. Descreva apenas fatos explicitamente fornecidos; não invente pessoas, emoções, intenção ou identidade. Responda em uma frase curta em português.")
        return try await sessao.respond(to: "Evidências visuais disponíveis: \(fatos). Gere uma observação editorial conservadora.").content
    } catch { return nil }
}

/// Processa um pedido por linha de stdin e escreve uma resposta por linha em stdout.
/// Mantém o processo vivo entre quadros para não recarregar Vision/Foundation Models a
/// cada frame — o custo de inicialização é o que tornava a análise lenta em lote.
@main struct Principal {
    static func main() async {
        while let linha = readLine(strippingNewline: true) {
            guard !linha.isEmpty, let entradaDados = linha.data(using: .utf8) else { continue }
            let saida: Saida
            do {
                let entrada = try JSONDecoder().decode(Entrada.self, from: entradaDados)
                var resultado = entrada.frames?.count ?? 0 > 1 ? analisarSequencia(entrada) : analisarVision(entrada)
                if entrada.resumir, let resumo = await interpretar(resultado) {
                    resultado = Saida(evidencias: resultado.evidencias + [Evidencia(
                        tipo: "interpretacao_editorial", valor: ["texto": .texto(resumo)],
                        confianca: nil, modelo: "AppleFoundationModels")], avisos: resultado.avisos)
                }
                saida = resultado
            } catch {
                saida = Saida(evidencias: [], avisos: ["erro_fatal: \(error.localizedDescription)"])
            }
            guard let dados = try? JSONEncoder().encode(saida), let linhaDeSaida = String(data: dados, encoding: .utf8) else {
                FileHandle.standardError.write(Data("Falha ao serializar resposta\n".utf8))
                continue
            }
            print(linhaDeSaida)
            fflush(stdout)
        }
    }
}